diff --git a/README.md b/README.md index 8b07ab2..f560657 100644 --- a/README.md +++ b/README.md @@ -2,13 +2,7 @@ LLM training in simple, pure C/CUDA. There is no need for 245MB of PyTorch or 107MB of cPython. Training GPT-2 (CPU, fp32) is ~1,000 lines of clean code in the single file [train_gpt2.c](train_gpt2.c), and training it on GPU is ~2,000 lines (adds CUDA kernels) in [train_gpt2.cu](train_gpt2.cu). The code compiles and runs instantly, it exactly matches the PyTorch reference implementation, and it ~matches the speed of (compiled) PyTorch (fp32, no flash attention). I chose GPT-2 as the first working example because it is the grand-daddy of LLMs, the first time the modern stack was put together. -Currently, we are working on: - -- optimize the CUDA implementation further to match/exceed PyTorch speed -- lower the precision from fp32 to mixed precision training -- add multi-gpu training, starting with DDP -- reproduce the GPT-2 training run (add data, evals) -- more modern architectures, Llama 2, Gemma, Mistral, etc. +Our current goal is to reproduce GPT-2 with a multi-node, mixed-precision, efficient implementation. For an overview of current ongoing work, see the latest [State of the Union](https://github.com/karpathy/llm.c/discussions/224) post. I'd like this repo to only maintain C and CUDA code. Ports of this repo to other languages are very welcome, but should be done in separate repos, and then I am happy to link to them below in the "notable forks" section, just like I did in [llama2.c notable forks](https://github.com/karpathy/llama2.c/tree/master?tab=readme-ov-file#notable-forks). @@ -255,10 +249,12 @@ Lastly, I will be a lot more sensitive to complexity in the root folder of the p - C# - [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): a C# port of this project - + - Rust - [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project +- Metal + - [llm.metal](https://github.com/regrettable-username/llm.metal) by @[regrettable-username](https://github.com/regrettable-username): LLM training in simple, raw C/Metal Shading Language ## discussions Ways of organizing development: diff --git a/dev/cuda/common.h b/dev/cuda/common.h index 6723a0f..cc3ba17 100644 --- a/dev/cuda/common.h +++ b/dev/cuda/common.h @@ -86,8 +86,8 @@ void validate_result(T* device_result, const T* cpu_reference, const char* name, if (i < 5) { printf("%f %f\n", cpu_reference[i], out_gpu[i]); } - // ensure correctness for all elements - if (fabs(cpu_reference[i] - out_gpu[i]) > tolerance) { + // ensure correctness for all elements. We can set an "ignore" mask by writing NaN + if (fabs(cpu_reference[i] - out_gpu[i]) > tolerance && !isnan(cpu_reference[i])) { printf("Mismatch of %s at %d: CPU_ref: %f vs GPU: %f\n", name, i, cpu_reference[i], out_gpu[i]); nfaults ++; if (nfaults >= 10) { diff --git a/dev/cuda/layernorm_backward.cu b/dev/cuda/layernorm_backward.cu index 1615123..132da59 100644 --- a/dev/cuda/layernorm_backward.cu +++ b/dev/cuda/layernorm_backward.cu @@ -6,6 +6,9 @@ nvcc -O3 --use_fast_math layernorm_backward.cu -o layernorm_backward version 1 is naive port from CPU code to kernel: parallelizes over B,T, loops over C ./layernorm_backward 1 + +version 2 moves a lot of reduction to shared memory over global memory +./layernorm_backward 2 */ #include @@ -152,19 +155,18 @@ __global__ void layernorm_backward_kernel1(float* dinp, float* dweight, float* d } } - -// super naive kernel that just parallelizes over B,T and loops over C +// uses shared memory instead for the reduces __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* dbias, const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd, int B, int T, int C) { + extern __shared__ float shared[]; // size = 2 * C + namespace cg = cooperative_groups; cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); int N = B * T; - if(idx >= N) { - return; - } + if(idx >= N) { return; } // thread guards int b = idx / T; int t = idx % T; @@ -175,6 +177,18 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d const float mean_bt = mean[b * T + t]; const float rstd_bt = rstd[b * T + t]; + // the first half of shared memory is bias, second is weight + float* dbias_shared = shared; + float* dweight_shared = shared + C; + + // init shared memory to zero + #pragma unroll + for(int i = threadIdx.x; i < C; i+= blockDim.x){ + dbias_shared[i] = 0.0f; + dweight_shared[i] = 0.0f; + } + __syncthreads(); + // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; @@ -184,10 +198,8 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } - dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus{}); dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus{}); - dnorm_mean = dnorm_mean / C; dnorm_norm_mean = dnorm_norm_mean / C; @@ -196,9 +208,9 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = weight[i] * dout_bt[i]; // gradient contribution to bias - atomicAdd(&dbias[i], dout_bt[i]); + atomicAdd(&dbias_shared[i], dout_bt[i]); // gradient contribution to weight - atomicAdd(&dweight[i], norm_bti * dout_bt[i]); + atomicAdd(&dweight_shared[i], norm_bti * dout_bt[i]); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 @@ -207,6 +219,13 @@ __global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* d dval *= rstd_bt; // final scale dinp_bt[i] += dval; } + __syncthreads(); + + // write to global memory + for(int i = threadIdx.x; i < C; i+= blockDim.x){ + atomicAdd(&dbias[i], dbias_shared[i]); + atomicAdd(&dweight[i], dweight_shared[i]); + } } // ---------------------------------------------------------------------------- @@ -225,7 +244,8 @@ void layernorm_backward2(float* dinp, float* dweight, float* dbias, int B, int T, int C, const int block_size) { const int N = B * T; const int grid_size = ceil_div(32*N, block_size); - layernorm_backward_kernel2<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); + size_t shared_mem_size = 2 * C * sizeof(float); + layernorm_backward_kernel2<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); } // kernel version dispatch @@ -358,4 +378,4 @@ int main(int argc, char **argv) { cudaCheck(cudaFree(d_rstd)); return 0; -} \ No newline at end of file +} diff --git a/dev/cuda/matmul_backward_bias.cu b/dev/cuda/matmul_backward_bias.cu index 15753d8..7feab39 100644 --- a/dev/cuda/matmul_backward_bias.cu +++ b/dev/cuda/matmul_backward_bias.cu @@ -7,6 +7,7 @@ nvcc -O3 matmul_backward_bias.cu -lineinfo -o matmul_backward_bias ./matmul_backward_bias 1 ./matmul_backward_bias 2 ./matmul_backward_bias 3 +./matmul_backward_bias 4 ncu: sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1 @@ -14,6 +15,7 @@ sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1 #include #include +#include #include #include #include @@ -124,6 +126,46 @@ __global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, in } } +// this kernel performs a column-wise reduction over dout, in PyTorch equivalent to: +// dbias = dout.sum((0,1)) +// the idea is to employ one block to reduce along several columns, +// where each block has a width of 32 columns to ensure coalesced access. +// at the end we accumulate the reductions performed by the warps in each block via shared memory +__global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, int B, int T, int OC) { + // this kernel is launched with 1D grid_dim of OC/32 + // for example let's say block_size is 128 + extern __shared__ float smem[]; // of size block_size (128) + const int warp_id = threadIdx.x / warpSize; // warp index in the block, 0,1,2,3 + const int lane_id = threadIdx.x % warpSize; // thread index in the warp, 0,1,2,...,31 + const int tl = blockIdx.x * warpSize; // pointer to the start column for this block + const int vstep = blockDim.x / warpSize; // number of warps in a block, e.g. 4 + + // pointer to the start of the column for one lane of threads + // so e.g. 4 threads (of the same lane_id) will reduce this one column + const float* dout_col = dout + tl + lane_id; + + // column reductions by looping through the rows + // each of the 4 threads offsets by its warp_id and then skips by vstep + // together these 4 threads cover all B*T rows of this (lane_id) column + // importantly, consecutive threads (in threadId) are processing adjacent columns, + // leading to a coalesced memory access pattern + float dout_sum = 0.0f; + for (int row = warp_id; row < B * T; row += vstep) { + dout_sum += dout_col[row * OC]; + } + smem[lane_id + warp_id * warpSize] = dout_sum; + __syncthreads(); + + // warp_id 0 reduces the shared memory column-wise, linearly + dout_sum = 0.0f; + if (warp_id == 0) { + for (int j = 0; j < vstep; j++) { + dout_sum += smem[lane_id + j * warpSize]; + } + dbias[tl + lane_id] += dout_sum; + } +} + // ---------------------------------------------------------------------------- // kernel launcher @@ -152,6 +194,14 @@ void matmul_backward_bias3(float* dinp, float* dweight, float* dbias, matmul_backward_bias_kernel3<<>>(dbias, dout, B, T, OC); } +void matmul_backward_bias4(float* dinp, float* dweight, float* dbias, + float* dout, float* inp, float* weight, float* ones, + int B, int T, int C, int OC, int block_size) { + assert(OC % 32 == 0); // OC must be divisible by 32 for this kernel + const int grid_size = OC / 32; + matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); +} + void matmul_backward_bias(int kernel_num, float* dinp, float* dweight, float* dbias, float* dout, float* inp, float* weight, float* ones, @@ -166,6 +216,9 @@ void matmul_backward_bias(int kernel_num, case 3: matmul_backward_bias3(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); break; + case 4: + matmul_backward_bias4(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); @@ -230,7 +283,7 @@ int main(int argc, char **argv) { matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, 128); // compare printf("Checking correctness...\n"); - validate_result(d_dbias, dbias, "dbias", OC, 1e-3f); + validate_result(d_dbias, dbias, "dbias", OC, 5e-3f); printf("All results match for block_size=%d.\n\n", block_size); } diff --git a/dev/cuda/trimat_forward.cu b/dev/cuda/trimat_forward.cu new file mode 100644 index 0000000..af8981a --- /dev/null +++ b/dev/cuda/trimat_forward.cu @@ -0,0 +1,596 @@ +/* +Triangular matrix multiplication as in autoregressive attention. A short story. +by @ngc92 + +Compile: +nvcc -O3 --use_fast_math trimat_forward.cu -o trimat_forward -lcublas + +Run: + +cuBLAS baseline kernel +./trimat_forward 0 + +naive +./trimat_forward 1 + +registers +./trimat_forward 2 + +tri3 +./trimat_forward 3 + +tri4 +./trimat_forward 4 +*/ + +#include +#include +#include +#include +#include +#include +#include +#include +#include "common.h" + +static cublasHandle_t cublas_handle; +static float* d_qkvr; // scratch for the cublas kernel + +/* ** Chapter I - Introduction ** + * + * You are Trimul. You've always wanted to do fast matrix multiplication, but they said + * "Don't bother, big dumb Cublas is much faster than you!" + * "I don't need to be faster than Cublas", you replied, "I can be smarter. Cublas calculates + * the entire matrix, but I need only half. If I'm more than half as fast as Cublas, I'm + * going to win." + * + * So to prove everyone wrong, you enter the TriMatlon, the most prestigious competition + * for anyone paying Attention. + * + * Before you start preparing, lets have a look at the players involved + * + * First, there is the Referee (`trimul_cpu`), slow and ponderous, but producing results + * beyond reproof. + * Then, there is Cublas. Cublas' mind is so inflexible, it doesn't actually comprehend + * what we are trying to do here, so Cublas has brought an assistant (`permute_kernel`) + * that translates the competition into a task that it can solve. But once it recognizes + * the problem, its muscle memory kicks in, and matrix products are produced faster than + * the eye can see. Stuck in its routine, Cublas doesn't realize the task is already + * finished with the lower triangle, though. + * + * If you can do without an assistant, and can solve the right task, then that's your opportunity + * to shine! + */ + + +// taken from then attention forward pass +void trimul_cpu(float* out, const float* inp, + int B, int T, int C, int NH) { + int C3 = C*3; + int hs = C / NH; // head size + float scale = 1.0 / sqrtf(hs); + + for (int b = 0; b < B; b++) { + for (int t = 0; t < T; t++) { + for (int h = 0; h < NH; h++) { + const float* query_t = inp + b * T * C3 + t * C3 + h * hs; + float* out_bth = out + b * NH * T * T + h * T * T + t * T; + + // pass 1: calculate query dot key and maxval + for (int t2 = 0; t2 <= t; t2++) { + const float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key + + // (query_t) dot (key_t2) + float val = 0.0f; + for (int i = 0; i < hs; i++) { + val += query_t[i] * key_t2[i]; + } + val *= scale; + + out_bth[t2] = val; + } + for(int t2 = t + 1; t2 < T; ++t2) { + out_bth[t2] = NAN; + } + } + } + } +} + +__global__ void permute_kernel(float* q, float* k, float* v, + const float* inp, + int B, int N, int NH, int d) { + // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d) + // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d) + int idx = blockIdx.x * blockDim.x + threadIdx.x; + + // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_] + + if (idx < B * NH * N * d) { + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + + int inp_idx = \ + (b * N * 3 * NH * d) + + (n * 3 * NH * d) + + (0 * NH * d) + + (nh_ * d) + + d_; + + q[idx] = inp[inp_idx]; + k[idx] = inp[inp_idx + NH * d]; + v[idx] = inp[inp_idx + 2 * (NH * d)]; + } +} + + +void trimul_cublas(float* preatt, + const float* inp, + int B, int T, int C, int NH) { + int HS = C / NH; // head size + + // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS) + float* q, * k, * v; + q = d_qkvr + 0 * B * T * C; + k = d_qkvr + 1 * B * T * C; + v = d_qkvr + 2 * B * T * C; + int total_threads = B * NH * T * HS; + int num_blocks = ceil_div(total_threads, 256); + permute_kernel<<>>(q, k, v, inp, B, T, NH, HS); + cudaCheck(cudaGetLastError()); + + // batched matrix multiply with cuBLAS + const float alpha = 1.0f / sqrtf(HS); + const float beta = 0.0f; + cublasCheck(cublasSgemmStridedBatched(cublas_handle, + CUBLAS_OP_T, CUBLAS_OP_N, + T, T, HS, + &alpha, + k, HS, T * HS, + q, HS, T * HS, + &beta, + preatt, T, T * T, + B * NH)); +} + +/* ** Chapter II - Getting a Team ** + * + * OK, you've registered for the competition, now what to do. TriMatlon is a team competition, so first, you need + * to figure out what kind of team you need, and how to organize it. The individual instances and heads of the + * problem are completely independent, so you just can send separate teams to work there completely independently. + * + * To figure out how to organize each team, you take out your spyglass (`Nsight Compute`) and look how the Cublas teams + * are handling their work. + * Turns out, you need 256 athletes in each group, and those handle 128 x 128 of the tasks. They work together in + * a tight square formation, 16 wide and 16 deep. + * + * So, you went out and got your 100 000 friends, and split them into groups (`trimul_launcher`). Each group gets + * informed about where they should work (`trimul_global`) and goes off to do their thing (`matmul_tri_naive`). + * Let's observe how we're doing. + */ + +// using creates an alias for a function pointer +using matmul_fn_ptr = void(*)(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha); + +template +__global__ void __launch_bounds__(256, 2) trimul_global(float* out, const float* inp, int T, int C, int NH) { + // skip above the diagonal + if(blockIdx.y > blockIdx.x) + return; + + // set up indices + int C3 = C*3; + int hs = C / NH; // head size + float scale = 1.0 / sqrtf(hs); + + // we put the "batch x head" dimension into the z block index. + int h = blockIdx.z % NH; + int b = blockIdx.z / NH; + + // Get the base address for the current batch and head + const float* q = inp + b * T * C3 + h * hs; + const float* k = inp + b * T * C3 + h * hs + C; + float* r = out + (b*NH + h)*T*T; + + // start the multiplication + matmul_tri(r, T, q, C3, k, C3, T, hs, scale); +} + +template +void trimul_launcher(float* out, const float* inp, int B, int T, int C, int NH) { + // we assume nice shapes here. Let's not make the code a mess by supporting weird shapes that you + // wouldn't want to use anyway. + assert(T % 128 == 0); + // No need to ceil_div, if it's not a multiple of 128, we would get wrong results anyway. + trimul_global<<>>(out, inp, T, C, NH); + cudaCheck(cudaGetLastError()); +} + +/* ** Chapter III - ... ** + * + * You go over to the playing field. On one end of the field, there is a huge pile of funnily shaped cookie cutters. + * Some in the shape of animals, some in the shape of a landscape. Each group of workers has assigned some runners, + * fetching the cookie cutters for them. The workers seem very relaxing, chatting with each other, lounging about. + * You focus in on one of them. + * + * He seems to be giving an instruction to a runner, and then turns back to reading a novel. The runner, meanwhile, + * crosses the field and back, handing him an elephant shape. Then she's off again to pick up a savannah background. + * Having received the two shapes, pressed them into the dough, and makes an elephant-in-the-savannah cookie. He hands + * the cutters back to the runner. "Can you please fetch me an elephant and a jungle next?" + * While she's on her way, he takes a sip off his cocktail. + * This time, she's making only one trip, keeping the elephant in her pocket (_Cache_). Still, it seems to take forever. + * You keep observing: + * - Elephant and zoo + * - Elephant and island + * ... + * - Lion and savannah + * - Lion and jungle + * - Lion and zoo + * ... + * + * The worker has his poor runner fetch the same things over and over again, looking like she's about to faint from exhaustion. + * Even though she realizes this and always keeps one of them in her pocket, there is so much running, + * and little actual work happening. + * + * Clearly, this isn't going to be effective, so you call a team meeting. + */ + +// baseline implementation: 20 ms +__device__ void matmul_tri_naive(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { + // get coordinates of our block + int i_base = 128 * blockIdx.x + 8 * threadIdx.x; + int j_base = 128 * blockIdx.y + 8 * threadIdx.y; + + // one more check to skip the upper diagonal in blocks that are on the diagonal. + if(j_base > i_base) + return; + + // Simple nested loop that calculates 8x8 results in one thread. + for(int io = 0; io < 8; ++io) { + int i = i_base + io; + for(int jo = 0; jo < 8; ++jo) { + int j = j_base + jo; + float val = 0; + for (int s = 0; s < hs; ++s) { + val += k[i * ks + s] * q[j * qs + s]; + } + p[i * ps + j] = val * alpha; + } + } +} + +/* ** Chapter IV - ... ** + * + * Each worker is producing 64 combined cookies from 8 animals and 8 landscapes. They send there runners of 64 times + * to fetch the corresponding shapes. This is terribly inefficient; The runners need a minute or so for each trip, + * but making a cookie can be done in just a second. + * + * "Let's try something different tomorrow: Just get all 16 cookie cutters that you need, and do all 64 combinations + * of them! See all this free space on your workbench (_registers_), you can keep them all there for easy access." + * + * The next morning, you come back to the field for another practice session. Initially, there is bustling activity + * with the runners, picking up 16 shapes for each worker. But then, the workers have to put down their newspapers + * and start making cookies. Now there are 64 combinations, so it takes them a full minute. + * + * Not all groups of workers are equally fast. When the first group finishes with all animal-landscape combinations, + * they already start asking the runners for the next set of cookie cutters, combining plants and houses. Even though + * the workers are much busier than before, they are still spending most of their time just waiting. + * + * Still, instead of being busy for 20 hours, your team is now done with the task in just 3h 30 minutes; already, this + * is five times faster. + * + * You think to yourself: "Why should we stop at 8 x 8 combinations? Lets to 16 x 16, that's only twice the work for + * the runners, but four times as much for the actual workers." + * You head over to the baking area, and make that suggestion to one of your team leaders. + * "In theory, that sounds great", she agrees, "but see, we only have limited space on our workbenches (_registers_). + * There is still some room left, but we simply cannot bake 256 cookies at the same time, sorry." + * + * A different strategy is needed, then. + */ + +// reorganize loops to enable data reuse: 3.5 ms +__device__ void matmul_tri_registers(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { + int i_base = 128 * blockIdx.x + 8 * threadIdx.x; + int j_base = 128 * blockIdx.y + 8 * threadIdx.y; + + if (j_base > i_base) + return; + + // shift our pointers to the sub-block this thread is responsible for + k += i_base * ks; + q += j_base * qs; + p += i_base * ps + j_base; + + float vals[8][8] = {}; + for (int s = 0; s < hs; ++s) { + float lhs[8]; + float rhs[8]; + for (int u = 0; u < 8; ++u) { + lhs[u] = k[u * ks + s]; + rhs[u] = q[u * qs + s]; + } + + for (int i = 0; i < 8; ++i) { + for (int j = 0; j < 8; ++j) { + vals[i][j] += lhs[i] * rhs[j]; + } + } + } + + for (int i = 0; i < 8; ++i) { + for (int j = 0; j < 8; ++j) { + p[i * ps + j] = vals[i][j] * alpha; + } + } +} + +/* ** Chapter IV - By the Bucketload ** + * + * Despite the hectic activity, you pick out one of the runners. "Why are you always brining just one shape? Wouldn't + * it be much more efficient if you took more than one?" + * "Of course", the runner answers, "but they've asked me for an elephant, a lion, a zebra, and a goldfish. These + * are all over the place, I can't just pick them up at one spot (_strided acccess_). + * "But the lion is right next to the palm tree. You could bring those two together?", you confirm. + * "Yes", he says, "if the just asked for the different categories at the same time, that would make things + * so much easier. See, I have this bucket, I could carry lots of things in one go if I could just scoop them up + * from the same place (_coalesced access_). + * + * OK, then lets fetch the first animal, first plant, first vehicle, and first landmark shape in one go (_vectorized load_). + * [Here, the metaphor breaks down a bit: Since we're accumulating all the results, getting more data at the same time + * depth-wise doesn't require more space on the workbench. We're stacking the cookies!] + * + * You also streamline the shape combination further. Instead of picking up all animals and landscapes at one, it is + * more efficient, using less workbench space, to just pick up all animals. Then, you get one landscape, combine it + * will all animals, get the next landscape, combine, and so on. + * + * In this way, instead of 2 x 8 x 4 cookie cutters that take up space, you only need (8+1) x 4 at the same time. + * + * With these optimizations, you are down to 100 minutes for this task. Still slower than Cublas, but not by much. + * + * In the arena, each team also has access to a small storage hut, much closer to their workbenches than the piles of + * cookie cutters on the other side. Cublas is using them heavily, so maybe you should, too. + */ + +// convenient helper functions to make the code below more readable +__device__ float4 ld_vec(const float* address) { + return *reinterpret_cast(address); +} + +__device__ void st_vec(float* address, float4 val) { + *reinterpret_cast(address) = val; +} + +// vector instructions for coalesced memory access: 1.7 ms +__device__ void matmul_tri3(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { + int i_base = 128 * blockIdx.x + 8 * threadIdx.x; + int j_base = 128 * blockIdx.y + 8 * threadIdx.y; + + if (j_base > i_base) + return; + + // shift our pointers to the sub-block this thread is responsible for + k += i_base * ks; + q += j_base * qs; + p += i_base * ps + j_base; + + float vals[8][8] = {}; + for (int s = 0; s < hs; s += 4) { + // load in float4 to improve coalescing + float4 rhs[8]; + for (int u = 0; u < 8; ++u) { + rhs[u] = ld_vec(q + u * qs + s); + } + + for (int i = 0; i < 8; ++i) { + // no need to keep lhs around for the i loop, its only reused in the j loop anyway. + float4 lhs = ld_vec(k + i * ks + s); + for (int j = 0; j < 8; ++j) { + vals[i][j] += lhs.x * rhs[j].x; + vals[i][j] += lhs.y * rhs[j].y; + vals[i][j] += lhs.z * rhs[j].z; + vals[i][j] += lhs.w * rhs[j].w; + } + } + } + + for (int i = 0; i < 8; ++i) { + for (int j = 0; j < 8; j += 4) { + float4 result; + result.x = vals[i][j + 0] * alpha; + result.y = vals[i][j + 1] * alpha; + result.z = vals[i][j + 2] * alpha; + result.w = vals[i][j + 3] * alpha; + st_vec(p + i * ps + j, result); + } + } +} + +/* ** Chapter V - Sharing is Caring ** + * + * You take a look around the shed, and see that there are 32 shelves there. They are much larger than the workbenches, + * giving you enough space for all the cookie cutters needed by the entire team. + * + * Within the team, workers have banded together in groups of 32. They are always doing the same thing, reducing the + * amount of effort required for coordination. However, that also means that if you send them all to pick up different + * cookie cutters from the same shelf, they will have to wait and queue up (_shared memory bank conflict_). + * + * In order to achieve maximum efficiency, we send the runners fetching cutters with the maximum bucket size: 32 different + * categories at the same time. + * + * [I'm having trouble getting the specifics into the story in a sensible way. For now, please read the code for more + * details.] + * + */ +__device__ void matmul_tri4(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { + int i_base = 128 * blockIdx.x + 8 * threadIdx.x; + int j_base = 128 * blockIdx.y + 8 * threadIdx.y; + + // we need all threads for loading data, so none of them can chicken out early, even + // if they are not responsible for any useful result. + if (blockIdx.y > blockIdx.x) + return; + + k += 128 * blockIdx.x * ks; + q += 128 * blockIdx.y * qs; + + __shared__ float lhs_s[128][32]; + __shared__ float rhs_s[128][32]; + + float vals[8][8] = {}; + for (int so = 0; so < hs; so += 32) { + // Read a large slice of the input, worked on together by all threads. + // They are organized differently for this part. We want to ensure + // fully coalesced loads, so we let a single warp handle consecutive + // addresses, which means we need to combine two threadIdx.y values + // in one read operation. + // note: threads may read data here that they don't need themselves. + // this really is a block-level operation. + __syncthreads(); + for(int y = threadIdx.y / 2; y < 128; y += 8) { + int xo = (threadIdx.y % 2) * 16; + lhs_s[y][threadIdx.x + xo] = k[y * ks + so + threadIdx.x + xo]; + rhs_s[y][threadIdx.x + xo] = q[y * qs + so + threadIdx.x + xo]; + } + __syncthreads(); + + for (int si = 0; si < 32; ++si) { + float rhs[8]; + for (int u = 0; u < 8; ++u) { + rhs[u] = rhs_s[u + 8 * threadIdx.y][(si + threadIdx.x) % 32]; + } + + for (int ii = 0; ii < 8; ++ii) { + float lhs = lhs_s[ii + 8 * threadIdx.x][(si + threadIdx.x) % 32]; + for (int ji = 0; ji < 8; ++ji) { + vals[ii][ji] += lhs * rhs[ji]; + } + } + } + } + + // don't write above the diagonal + if (j_base > i_base) + return; + + for (int ii = 0; ii < 8; ++ii) { + for (int ji = 0; ji < 8; ji += 4) { + int i = i_base + ii; + int j = j_base + ji; + float4 result; + result.x = vals[ii][ji + 0] * alpha; + result.y = vals[ii][ji + 1] * alpha; + result.z = vals[ii][ji + 2] * alpha; + result.w = vals[ii][ji + 3] * alpha; + st_vec(p + i * ps + j, result); + } + } +} + +/* ** Chapter VI - Competition Day ** + * + * Finally, you feel ready to take on Cublas. You hand out tickets to the event for you friends to see. + * + * --------------------------------------------------------------------------------- + * | CuBLAS vs TriMul - Fight of the Century | + * | | + * | Ticket code: | + * | > nvcc -O3 --use_fast_math trimat_forward.cu -o trimat_forward -lcublas | + * | > ./trimat 4 | + * | | + * --------------------------------------------------------------------------------- + */ + +void trimul_gpu(int kernel_num, + float* out, const float* inp, + int B, int T, int C, int NH) { + switch (kernel_num) { + case 0: + trimul_cublas(out, inp, B, T, C, NH); + break; + case 1: + trimul_launcher(out, inp, B, T, C, NH); + break; + case 2: + trimul_launcher(out, inp, B, T, C, NH); + break; + case 3: + trimul_launcher(out, inp, B, T, C, NH); + break; + case 4: + trimul_launcher(out, inp, B, T, C, NH); + break; + default: + printf("Invalid kernel number\n"); + exit(1); + } +} + + + +int main(int argc, char **argv) { + srand(0); + + int B = 8; + int T = 1024; + int C = 768; + int NH = 12; + + int deviceIdx = 0; + cudaCheck(cudaSetDevice(deviceIdx)); + cublasCreate(&cublas_handle); + + // create host memory of random numbers + float* out = (float*)malloc(B * NH * T * T * sizeof(float)); + float* inp = make_random_float(B * T * 3 * C); + + // move to GPU + float* d_out; + float* d_inp; + cudaCheck(cudaMalloc(&d_out, B * NH * T * T * sizeof(float))); + cudaCheck(cudaMalloc(&d_inp, B * T * 3 * C * sizeof(float))); + cudaCheck(cudaMemcpy(d_inp, inp, B * T * 3 * C * sizeof(float), cudaMemcpyHostToDevice)); + + // buffer for cublas + cudaCheck(cudaMalloc(&d_qkvr, B * T * 3 * C * sizeof(float))); + + // read kernel_num from command line + int kernel_num = 1; + if (argc > 1) { + kernel_num = atoi(argv[1]); + } + printf("Using kernel %d\n", kernel_num); + + // first check the correctness of the kernel + trimul_cpu(out, inp, B, T, C, NH); + trimul_gpu(kernel_num, d_out, d_inp, B, T, C, NH); + validate_result(d_out, out, "out", B * NH * T * T, 1e-4f); + + printf("All results match. Starting benchmarks.\n\n"); + + // benchmark speed of the kernel + int repeat_times = 100; + + float elapsed_time = benchmark_kernel(repeat_times, trimul_gpu, + kernel_num, d_out, d_inp, + B, T, C, NH); + + + float cublas_time = benchmark_kernel(repeat_times, trimul_gpu, + 0, d_out, d_inp, + B, T, C, NH); + + printf("time %.2f ms vs %.2f ms for CuBLAS\n", elapsed_time, cublas_time); + + // free memory + free(out); + free(inp); + cudaCheck(cudaFree(d_out)); + cudaCheck(cudaFree(d_inp)); + cublasDestroy(cublas_handle); + + return 0; +} diff --git a/train_gpt2.c b/train_gpt2.c index 87a44bf..6028d1a 100644 --- a/train_gpt2.c +++ b/train_gpt2.c @@ -364,8 +364,10 @@ void gelu_forward(float* out, float* inp, int N) { } // we want to use -Ofast optimization, but sadly GeLU breaks, so disable this flag just for it (#168) -#pragma float_control(precise, on, push) // On msvc /fp:fast is a lot faster, but the expf inside coshf breaks the model -__attribute__((optimize("no-finite-math-only"))) // same for gcc -Ofast +#pragma float_control(precise, on, push) +#if defined(__GNUC__) && !defined(__clang__) +__attribute__((optimize("no-finite-math-only"))) +#endif void gelu_backward(float* dinp, float* inp, float* dout, int N) { for (int i = 0; i < N; i++) { float x = inp[i]; diff --git a/train_gpt2.cu b/train_gpt2.cu index e4f9c49..f1008f6 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -664,51 +664,79 @@ __global__ void softmax_forward_kernel7(float* out, const float* inp, int N, int } } -// cooperative groups solution, one warp per output channel -template -__global__ void matmul_backward_bias_kernel2(Td* dbias, const Td* dout, int B, int T, int OC) { - // dout is (B, T, OC), dbias is (OC) - // e.g. if block_size = 128, then we have 4 warps per block, each in charge of one output channel - namespace cg = cooperative_groups; - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); - // meta_group_size is the number of warps in a block (e.g. 4), meta_group_rank is the warp index (0,1,2,3) - int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); - if(idx >= OC) { return; } - int BT = B * T; // number of elements to reduce in total, per channel - // first, thread coarsening to sum reduce the problem size from B*T to 32 - float sum = 0.0f; - for(int i = warp.thread_rank(); i < BT; i += warp.size()) { - sum += (float)dout[i * OC + idx]; +// this kernel performs a column-wise reduction over dout, in PyTorch equivalent to: +// dbias = dout.sum((0,1)) +// the idea is to employ one block to reduce along several columns, +// where each block has a width of 32 columns to ensure coalesced access. +// at the end we accumulate the reductions performed by the warps in each block via shared memory +__global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, int B, int T, int OC) { + // this kernel is launched with 1D grid_dim of OC/32 + // for example let's say block_size is 128 + extern __shared__ float smem[]; // of size block_size (128) + const int warp_id = threadIdx.x / warpSize; // warp index in the block, 0,1,2,3 + const int lane_id = threadIdx.x % warpSize; // thread index in the warp, 0,1,2,...,31 + const int tl = blockIdx.x * warpSize; // pointer to the start column for this block + const int vstep = blockDim.x / warpSize; // number of warps in a block, e.g. 4 + + // pointer to the start of the column for one lane of threads + // so e.g. 4 threads (of the same lane_id) will reduce this one column + const float* dout_col = dout + tl + lane_id; + + // column reductions by looping through the rows + // each of the 4 threads offsets by its warp_id and then skips by vstep + // together these 4 threads cover all B*T rows of this (lane_id) column + // importantly, consecutive threads (in threadId) are processing adjacent columns, + // leading to a coalesced memory access pattern + float dout_sum = 0.0f; + for (int row = warp_id; row < B * T; row += vstep) { + dout_sum += dout_col[row * OC]; } - // now do a warp-level reduce to get the sum across the 32 threads in this warp - sum = cg::reduce(warp, sum, cg::plus{}); - // write the result to output (global memory) - if(warp.thread_rank() == 0) { - dbias[idx] = (Td)((float)dbias[idx] + sum); + smem[lane_id + warp_id * warpSize] = dout_sum; + __syncthreads(); + + // warp_id 0 reduces the shared memory column-wise, linearly + dout_sum = 0.0f; + if (warp_id == 0) { + for (int j = 0; j < vstep; j++) { + dout_sum += smem[lane_id + j * warpSize]; + } + dbias[tl + lane_id] += dout_sum; } } -template -__global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams* dbias, - const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, - int B, int T, int C) { +// uses shared memory instead for the reduces +__global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* dbias, + const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd, + int B, int T, int C) { + extern __shared__ float shared[]; // size = 2 * C + + namespace cg = cooperative_groups; cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); int N = B * T; - if(idx >= N) { - return; - } + if(idx >= N) { return; } // thread guards int b = idx / T; int t = idx % T; - const Tdout* dout_bt = dout + b * T * C + t * C; - const Trest* inp_bt = inp + b * T * C + t * C; - Tdinp* dinp_bt = dinp + b * T * C + t * C; - float mean_bt = (float)mean[b * T + t]; - float rstd_bt = (float)rstd[b * T + t]; + const float* dout_bt = dout + b * T * C + t * C; + const float* inp_bt = inp + b * T * C + t * C; + float* dinp_bt = dinp + b * T * C + t * C; + const float mean_bt = mean[b * T + t]; + const float rstd_bt = rstd[b * T + t]; + + // the first half of shared memory is bias, second is weight + float* dbias_shared = shared; + float* dweight_shared = shared + C; + + // init shared memory to zero + #pragma unroll + for(int i = threadIdx.x; i < C; i+= blockDim.x){ + dbias_shared[i] = 0.0f; + dweight_shared[i] = 0.0f; + } + __syncthreads(); // first: two reduce operations float dnorm_mean = 0.0f; @@ -719,10 +747,8 @@ __global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } - dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus{}); dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus{}); - dnorm_mean = dnorm_mean / C; dnorm_norm_mean = dnorm_norm_mean / C; @@ -731,9 +757,9 @@ __global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; // gradient contribution to bias - atomicAddX(&dbias[i], (Tparams)dout_bt[i]); + atomicAdd(&dbias_shared[i], dout_bt[i]); // gradient contribution to weight - atomicAddX(&dweight[i], (Tparams)(norm_bti * (float)dout_bt[i])); + atomicAdd(&dweight_shared[i], norm_bti * dout_bt[i]); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 @@ -742,6 +768,13 @@ __global__ void layernorm_backward_kernel(Tdinp* dinp, Tparams* dweight, Tparams dval *= rstd_bt; // final scale dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval); } + __syncthreads(); + + // write to global memory + for(int i = threadIdx.x; i < C; i+= blockDim.x){ + atomicAdd(&dbias[i], dbias_shared[i]); + atomicAdd(&dweight[i], dweight_shared[i]); + } } @@ -1164,9 +1197,9 @@ void matmul_backward_fp32(float* dinp, float* dweight, float* dbias, cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one, inp, C, dout, OC, &one, dweight, C)); // backward to bias, if given, does a += if (dbias != NULL) { - const int block_size = 512; - const int grid_size = CEIL_DIV(OC * 32, block_size); - matmul_backward_bias_kernel2<<>>(dbias, dout, B, T, OC); + const int block_size = 1024; + const int grid_size = OC / 32; // for now, OC must be divisible by 32 for this kernel to work + matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); cudaCheck(cudaGetLastError()); } } @@ -1193,11 +1226,11 @@ template void layernorm_backward(Tdinp* dinp, Tparams* dweight, Tparams* dbias, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C) { - const int block_size = 256; + const int block_size = 512; const int N = B * T; - // one warp per token, so we need to divide by 32 here. - const int grid_size = CEIL_DIV(N, block_size / 32); - layernorm_backward_kernel<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); + const int grid_size = CEIL_DIV(32*N, block_size); + size_t shared_mem_size = 2 * C * sizeof(float); + layernorm_backward_kernel2<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); cudaCheck(cudaGetLastError()); }