diff --git a/llmc/encoder.cuh b/llmc/encoder.cuh new file mode 100644 index 0000000..67e37bc --- /dev/null +++ b/llmc/encoder.cuh @@ -0,0 +1,131 @@ +#include "cuda_common.h" +#include "cuda_utils.cuh" + +__global__ void encoder_forward_kernel3(floatX* out, + const int* inp, const floatX* wte, const floatX* wpe, + int B, int T, int C) { + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; + int N = B * T * C; + if (idx >= N) { return; } + + int bt = idx / C; + int b = bt / T; + int t = bt % T; + int c = idx % C; + + int ix = inp[b * T + t]; + + floatX* out_btc = out + b * T * C + t * C + c; + const floatX* wte_ix = wte + ix * C + c; + const floatX* wpe_tc = wpe + t * C + c; + + x128 packed_out; + x128 wte128 = load128cs(wte_ix); + x128 wpe128 = load128cs(wpe_tc); + for (int k = 0; k < x128::size; k++) { + packed_out[k] = (floatX)((float)wte128[k] + (float)wpe128[k]); + } + store128(out_btc, packed_out); +} + +template +__global__ void wte_backward_kernel(floatX* dwte, + const int4* bucket_info, const int* workload_indices, const floatX* dout, const int* inp, + unsigned int seed, int B, int T, int C) { + // In order to be deterministic, we preprocess the inputs on the cpu into "buckets" + // Each bucket corresponds to (WARP_SIZE * x128::size) channels for a single vocabulary token + // Each thread handles x128::size channels, e.g. 256 per warp for BF16 + // Each block handles (BLOCK_SIZE / WARP_SIZE) elements in a single bucket in parallel + // If a bucket has less than 8 elements, some warps will return immediately + // If a bucket has more than 8 elements, we will loop over all of them + // The buckets are sorted on the CPU so the largest buckets start 1st + int bucket = blockIdx.x; + int warp_id = threadIdx.x / WARP_SIZE; + int lane_id = threadIdx.x % WARP_SIZE; + int c_per_warp = WARP_SIZE * x128::size; + + int bucket_start_idx = bucket_info[bucket].x; + int bucket_size = bucket_info[bucket].y; + int bucket_ix = bucket_info[bucket].z; + int c = bucket_info[bucket].w * c_per_warp + (lane_id * x128::size); + + // Each thread handles "x128::size" channels, so at fp8, each warp would handle 512 channels + // If C is not a multiple of this (e.g. 768), some buckets/c_groups cannot use the entire warp + if (c >= C) { return; } + // Exit early if this is a small bucket and this warp doesn't have any items to process + if (warp_id >= bucket_size) { return; } + + float accum[x128::size] = {0.0f}; + __shared__ float accum_shared[x128::size * BLOCK_SIZE]; + + for(int item = warp_id; item < bucket_size; item += BLOCK_SIZE/WARP_SIZE) { + int bt = workload_indices[bucket_start_idx + item]; + + const floatX* dout_btc = dout + bt * C + c; + x128 packed_inp1 = load128cs(dout_btc); + for (int k = 0; k < packed_inp1.size; k++) { + accum[k] += (float)packed_inp1[k]; + } + } + + if (warp_id != 0) { + // we accumulate into warp 0, so only the other warps need to write to shared memory + for (int k = 0; k < x128::size; k++) { + accum_shared[threadIdx.x + k * BLOCK_SIZE] = accum[k]; + } + return; // only warp 0 is needed after writing to shared memory + } + + // Read dwte for warp 0 even if other warps are not finished yet to maximise latency tolerance + floatX* dwte_ix = dwte + bucket_ix * C + c; + x128 packed_in_out = load128(dwte_ix); + + // note: threads which have returned are considered synchronised by CUDA so no risk of deadlock + __syncthreads(); + + // Accumulate into warp 0's registers by reading the values of the other warps in shared memory + for (int i = threadIdx.x+WARP_SIZE; i < min(BLOCK_SIZE, bucket_size*WARP_SIZE); i += WARP_SIZE) { + for (int k = 0; k < x128::size; k++) { + accum[k] += accum_shared[i + k * BLOCK_SIZE]; + } + } + + // Add the result to dwte and write back to global memory (read-modify-write) + for (unsigned int k = 0; k < x128::size; k++) { + // We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic + stochastic_rounding(accum[k] + (float)packed_in_out[k], &packed_in_out[k], seed + k); + } + store128(dwte_ix, packed_in_out); +} + +__global__ void wpe_backward_kernel(floatX* dwpe, + const floatX* dout, const int* inp, + int B, int T, int C, unsigned int seed) { + // Each thread handles x128::size "channel positions", e.g. 256 per warp for BF16 + // For gpt2-124M BF16, C=768 and T=1024, so 3 warps per channel and 3072 warps in total + // For each "channel position" we sum the gradients for every batch at that C/T element + // This way each dwte element is only updated once, and the kernel is fully deterministic! + // The previous kernel was not deterministic, as batches were aggregated with atomicAdd + int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; + if (idx >= T * C) { return; } + + // if C is not a multiple of WARP_SIZE*x128::size, it's OK for some warps to handle multiple t + int t = idx / C; + int c = idx % C; + float accum[x128::size] = {0.0f}; + + for (int b = 0; b < B; b++) { + x128 packed_dout = load128cs(dout + (b * T * C) + (t * C) + c); // will never be read again + for (int k = 0; k < x128::size; k++) { + accum[k] += (float)packed_dout[k]; + } + } + + floatX* dwpe_tc = dwpe + (t * C) + c; + x128 packed_dwpe = load128(dwpe_tc); + for (unsigned int k = 0; k < x128::size; k++) { + // We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic + stochastic_rounding(accum[k] + (float)packed_dwpe[k], &packed_dwpe[k], seed + k); + } + store128(dwpe_tc, packed_dwpe); +} diff --git a/train_gpt2.cu b/train_gpt2.cu index 0c29d53..ee66ebb 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -46,6 +46,8 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. // defines: CUBLAS_LOWP, cublasCheck, cublaslt_workspace_size, cublaslt_workspace // defines: cublas_compute, cublaslt_handle, cublas_handle #include "llmc/cublas_common.h" +// defines: encoder_forward_kernel3, wte_backward_kernel, wpe_backward_kernel +#include "llmc/encoder.cuh" // defines: create_cudnn, destroy_cudnn, attention_forward_cudnn, attention_backward_cudnn #ifdef ENABLE_CUDNN #include "llmc/cudnn_att.h" @@ -245,135 +247,6 @@ gradients have to add. We make sure that we do a += as necessary. E.g., the layernorms are connected to the residuals so we += in layernorm backward. */ -__global__ void encoder_forward_kernel3(floatX* out, - const int* inp, const floatX* wte, const floatX* wpe, - int B, int T, int C) { - int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; - int N = B * T * C; - if (idx >= N) { return; } - - int bt = idx / C; - int b = bt / T; - int t = bt % T; - int c = idx % C; - - int ix = inp[b * T + t]; - - floatX* out_btc = out + b * T * C + t * C + c; - const floatX* wte_ix = wte + ix * C + c; - const floatX* wpe_tc = wpe + t * C + c; - - x128 packed_out; - x128 wte128 = load128cs(wte_ix); - x128 wpe128 = load128cs(wpe_tc); - for (int k = 0; k < x128::size; k++) { - packed_out[k] = (floatX)((float)wte128[k] + (float)wpe128[k]); - } - store128(out_btc, packed_out); -} - -template -__global__ void wte_backward_kernel(floatX* dwte, - const int4* bucket_info, const int* workload_indices, const floatX* dout, const int* inp, - unsigned int seed, int B, int T, int C) { - // In order to be deterministic, we preprocess the inputs on the cpu into "buckets" - // Each bucket corresponds to (WARP_SIZE * x128::size) channels for a single vocabulary token - // Each thread handles x128::size channels, e.g. 256 per warp for BF16 - // Each block handles (BLOCK_SIZE / WARP_SIZE) elements in a single bucket in parallel - // If a bucket has less than 8 elements, some warps will return immediately - // If a bucket has more than 8 elements, we will loop over all of them - // The buckets are sorted on the CPU so the largest buckets start 1st - int bucket = blockIdx.x; - int warp_id = threadIdx.x / WARP_SIZE; - int lane_id = threadIdx.x % WARP_SIZE; - int c_per_warp = WARP_SIZE * x128::size; - - int bucket_start_idx = bucket_info[bucket].x; - int bucket_size = bucket_info[bucket].y; - int bucket_ix = bucket_info[bucket].z; - int c = bucket_info[bucket].w * c_per_warp + (lane_id * x128::size); - - // Each thread handles "x128::size" channels, so at fp8, each warp would handle 512 channels - // If C is not a multiple of this (e.g. 768), some buckets/c_groups cannot use the entire warp - if (c >= C) { return; } - // Exit early if this is a small bucket and this warp doesn't have any items to process - if (warp_id >= bucket_size) { return; } - - float accum[x128::size] = {0.0f}; - __shared__ float accum_shared[x128::size * BLOCK_SIZE]; - - for(int item = warp_id; item < bucket_size; item += BLOCK_SIZE/WARP_SIZE) { - int bt = workload_indices[bucket_start_idx + item]; - - const floatX* dout_btc = dout + bt * C + c; - x128 packed_inp1 = load128cs(dout_btc); - for (int k = 0; k < packed_inp1.size; k++) { - accum[k] += (float)packed_inp1[k]; - } - } - - if (warp_id != 0) { - // we accumulate into warp 0, so only the other warps need to write to shared memory - for (int k = 0; k < x128::size; k++) { - accum_shared[threadIdx.x + k * BLOCK_SIZE] = accum[k]; - } - return; // only warp 0 is needed after writing to shared memory - } - - // Read dwte for warp 0 even if other warps are not finished yet to maximise latency tolerance - floatX* dwte_ix = dwte + bucket_ix * C + c; - x128 packed_in_out = load128(dwte_ix); - - // note: threads which have returned are considered synchronised by CUDA so no risk of deadlock - __syncthreads(); - - // Accumulate into warp 0's registers by reading the values of the other warps in shared memory - for (int i = threadIdx.x+WARP_SIZE; i < min(BLOCK_SIZE, bucket_size*WARP_SIZE); i += WARP_SIZE) { - for (int k = 0; k < x128::size; k++) { - accum[k] += accum_shared[i + k * BLOCK_SIZE]; - } - } - - // Add the result to dwte and write back to global memory (read-modify-write) - for (unsigned int k = 0; k < x128::size; k++) { - // We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic - stochastic_rounding(accum[k] + (float)packed_in_out[k], &packed_in_out[k], seed + k); - } - store128(dwte_ix, packed_in_out); -} - -__global__ void wpe_backward_kernel(floatX* dwpe, - const floatX* dout, const int* inp, - int B, int T, int C, unsigned int seed) { - // Each thread handles x128::size "channel positions", e.g. 256 per warp for BF16 - // For gpt2-124M BF16, C=768 and T=1024, so 3 warps per channel and 3072 warps in total - // For each "channel position" we sum the gradients for every batch at that C/T element - // This way each dwte element is only updated once, and the kernel is fully deterministic! - // The previous kernel was not deterministic, as batches were aggregated with atomicAdd - int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; - if (idx >= T * C) { return; } - - // if C is not a multiple of WARP_SIZE*x128::size, it's OK for some warps to handle multiple t - int t = idx / C; - int c = idx % C; - float accum[x128::size] = {0.0f}; - - for (int b = 0; b < B; b++) { - x128 packed_dout = load128cs(dout + (b * T * C) + (t * C) + c); // will never be read again - for (int k = 0; k < x128::size; k++) { - accum[k] += (float)packed_dout[k]; - } - } - - floatX* dwpe_tc = dwpe + (t * C) + c; - x128 packed_dwpe = load128(dwpe_tc); - for (unsigned int k = 0; k < x128::size; k++) { - // We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic - stochastic_rounding(accum[k] + (float)packed_dwpe[k], &packed_dwpe[k], seed + k); - } - store128(dwpe_tc, packed_dwpe); -} - __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, const floatX* __restrict__ inp, const floatX* __restrict__ weight, const floatX* __restrict__ bias, int N, int C) {