From b58181079ae2e7c770871bb651dbdfba46441c2f Mon Sep 17 00:00:00 2001 From: ademeure Date: Fri, 3 May 2024 16:38:46 +0100 Subject: [PATCH] super optimised matmul_bias + cuda streams for better parallelism + fix fused_classifier cache hint + remove remaining cooperative groups --- profile_gpt2.cu | 15 ++- test_gpt2.cu | 16 ++- train_gpt2.cu | 290 +++++++++++++++++++++++++++++++----------------- 3 files changed, 217 insertions(+), 104 deletions(-) diff --git a/profile_gpt2.cu b/profile_gpt2.cu index ed09447..ce00f7c 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -34,11 +34,22 @@ int main() { cudaCheck(cudaSetDevice(deviceIdx)); cudaDeviceProp deviceProp; cudaGetDeviceProperties(&deviceProp, deviceIdx); - cuda_num_SMs = deviceProp.multiProcessorCount; - cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor; printf("[System]\n"); printf("Device %d: %s\n", deviceIdx, deviceProp.name); + cuda_num_SMs = deviceProp.multiProcessorCount; + cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor; + cuda_arch_major = deviceProp.major; + cuda_arch_minor = deviceProp.minor; + + cudaCheck(cudaStreamCreate(&main_stream)); + cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming); + cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming); + for (int i = 0; i < num_parallel_streams; i++) { + cudaCheck(cudaStreamCreate(¶llel_streams[i])); + cudaEventCreateWithFlags(¶llel_events[i], cudaEventDisableTiming); + } + // setup cuBLAS and cuBLASLt cublasCheck(cublasCreate(&cublas_handle)); cublasCheck(cublasLtCreate(&cublaslt_handle)); diff --git a/test_gpt2.cu b/test_gpt2.cu index 67afa50..34ab621 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -89,12 +89,22 @@ int main(int argc, char *argv[]) { cudaCheck(cudaSetDevice(deviceIdx)); cudaDeviceProp deviceProp; cudaGetDeviceProperties(&deviceProp, deviceIdx); - cuda_num_SMs = deviceProp.multiProcessorCount; - cuda_arch_major = deviceProp.major; - cuda_arch_minor = deviceProp.minor; printf("[System]\n"); printf("Device %d: %s\n", deviceIdx, deviceProp.name); + cuda_num_SMs = deviceProp.multiProcessorCount; + cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor; + cuda_arch_major = deviceProp.major; + cuda_arch_minor = deviceProp.minor; + + cudaCheck(cudaStreamCreate(&main_stream)); + cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming); + cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming); + for (int i = 0; i < num_parallel_streams; i++) { + cudaCheck(cudaStreamCreate(¶llel_streams[i])); + cudaEventCreateWithFlags(¶llel_events[i], cudaEventDisableTiming); + } + // setup cuBLAS and cuBLASLt cublasCheck(cublasCreate(&cublas_handle)); cublasCheck(cublasLtCreate(&cublaslt_handle)); diff --git a/train_gpt2.cu b/train_gpt2.cu index 9c48667..100f988 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -51,8 +51,6 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200), #include #include #include -#include -#include #include // Multi-GPU related @@ -156,11 +154,13 @@ int cuda_arch_minor = 0; int cuda_num_SMs = 0; // for persistent threads where we want 1 threadblock per SM int cuda_threads_per_SM = 0; -constexpr int num_parallel_streams = 4; // + 1 primary "stream" (+ default stream) +// CUDA streams & events (note: non-timing events, use separate event for timing/profiling!) +constexpr int num_parallel_streams = 2; // + 1 primary "main_stream" (+ default stream) cudaStream_t parallel_streams[num_parallel_streams]; -cudaStream_t stream; - -namespace cg = cooperative_groups; +cudaEvent_t parallel_events[num_parallel_streams]; +cudaStream_t main_stream; +cudaEvent_t main_event; +cudaEvent_t loss_event; // to make sure fused_classifier has written the losses to the CPU buffer // convenience macro for calculating grid/block dimensions for kernels #define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) @@ -261,7 +261,9 @@ __device__ floatX warpReduceSum(floatX val) { #endif using reduction_func_t = float (*) (float); -__device__ float blockReduce(float val, reduction_func_t warp_reduction) { +__device__ float blockReduce(float val, reduction_func_t warp_reduction, bool final_sync=false) { + // two reductions of up to 1024 threads: + // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) __shared__ float shared_val[32]; int lane_id = threadIdx.x % 32; int warp_id = threadIdx.x / 32; @@ -273,6 +275,10 @@ __device__ float blockReduce(float val, reduction_func_t warp_reduction) { // same strategy, now reduce across warps warp_val = (lane_id < num_warps) ? shared_val[lane_id] : 0.0f; float block_val = warp_reduction(warp_val); + + if (final_sync) { + __syncthreads(); // only needed in loops when effectively reusing shared memory etc. + } return block_val; } @@ -828,9 +834,12 @@ __global__ void encoder_backward_kernel(floatX* dwte, floatX* 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) { - 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(); + const int warp_size = 32; + int lane_id = threadIdx.x % warp_size; + int warp_id = threadIdx.x / warp_size; + int num_warps = blockDim.x / warp_size; + + int idx = blockIdx.x * num_warps + warp_id; if(idx >= N) { return; } // guard // the row of input that this group of threads is responsible for @@ -838,30 +847,30 @@ __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __re // mean float sum = 0.0f; - for (int i = warp.thread_rank(); i < C; i += warp.size()) { + for (int i = lane_id; i < C; i += warp_size) { sum += (float)x[i]; } - sum = cg::reduce(warp, sum, cg::plus{}); + sum = warpReduceSum(sum); float m = sum / C; - if(warp.thread_rank() == 0 && mean != nullptr) { + if(lane_id == 0 && mean != nullptr) { __stcs(mean + idx, (floatX)m); } // rstd sum = 0.0f; - for (int i = warp.thread_rank(); i < C; i += warp.size()) { + for (int i = lane_id; i < C; i += warp_size) { float diff = (float)x[i] - m; sum += diff * diff; } - sum = cg::reduce(warp, sum, cg::plus{}); + sum = warpReduceSum(sum); float s = rsqrtf(sum / C + 1e-5f); - if(warp.thread_rank() == 0 && rstd != nullptr) { + if(lane_id == 0 && rstd != nullptr) { __stcs(rstd + idx, (floatX)s); } // final normalization and scaling by weight/bias floatX* o = out + idx * C; - for (int c = warp.thread_rank(); c < C; c += warp.size()) { + for (int c = lane_id; c < C; c += warp_size) { // load and store using the .cs "streaming" hint to the compiler, // indicating that this data will not be reused soon, and can be streamed through the caches // this allows the threads to get more cache-hits for the (shared) weight and bias parameters @@ -948,14 +957,17 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons // directly autoregressive, so we only compute the lower triangular part // uses the online softmax algorithm assert(T % 4 == 0); - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + const int warp_size = 32; + int lane_id = threadIdx.x % warp_size; + int warp_id = threadIdx.x / warp_size; + int num_warps = blockDim.x / warp_size; + // micro-optimization: we iterate backwards so that // after the softmax backward operation completes, the cache retains the // part of the matrix close to the upper left corner, which benefits the // matmul operation that immediately follows. // int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // forward order - int idx = (gridDim.x - blockIdx.x - 1) * warp.meta_group_size() + warp.meta_group_rank(); // backward order + int idx = (gridDim.x - blockIdx.x - 1) * num_warps + warp_id; // backward order if(idx >= N * T) { return; } @@ -970,7 +982,7 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons float sumval = 0.0f; const floatX* x_aligned = reinterpret_cast(__builtin_assume_aligned(x, 16)); - for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) { + for (int i = lane_id; i < pos_by_4; i += warp_size) { float regarray[4]; #pragma unroll for (int k = 0; k < 4; ++k) { @@ -986,21 +998,21 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons } } - if(4*pos_by_4 + warp.thread_rank() <= own_pos) { + if(4*pos_by_4 + lane_id <= own_pos) { float old_maxval = maxval; - maxval = fmaxf(maxval, (float)x[4*pos_by_4 + warp.thread_rank()]); + maxval = fmaxf(maxval, (float)x[4*pos_by_4 + lane_id]); sumval *= expf(inv_temperature * (old_maxval - maxval)); - sumval += expf(inv_temperature * ((float)x[4*pos_by_4 + warp.thread_rank()] - maxval)); + sumval += expf(inv_temperature * ((float)x[4*pos_by_4 + lane_id] - maxval)); } - float global_maxval = cg::reduce(warp, maxval, cg::greater{}); + float global_maxval = warpReduceMax(maxval); sumval *= expf(inv_temperature * (maxval - global_maxval)); - float sum = cg::reduce(warp, sumval, cg::plus{}); + float sum = warpReduceSum(sumval); float norm = 1.f / sum; // divide the whole row by the sum - for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { + for (int i = lane_id; i <= own_pos; i += warp_size) { // recalculation is faster than doing the round-trip through memory. float ev = expf(inv_temperature * ((float)__ldcs(x + i) - global_maxval)); __stcs(out + idx * T + i, (floatX)(ev * norm)); @@ -1059,19 +1071,53 @@ __global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floa } } -__global__ void matmul_backward_bias_kernel5(float* dbias, const floatX* dout, int B, int T, int OC) { +__global__ void matmul_backward_bias_kernel6(float* dbias, const floatX* dout, int B, int T, int OC) { // note: this kernel reads in floatX, but it writes to float! // this is because we're using atomics, which are super slow in < fp32 precision on < H100 GPUs // so the trick is do fp32 atomics to a buffer, and then copy_and_cast the result to floatX - int oc = blockIdx.x * blockDim.x + threadIdx.x; - if(oc >= OC) return; - float sum = 0.0; - // grid-wide loop for maximum parallelism - for (int i = blockIdx.y; i < B * T; i += gridDim.y) { - sum += (float)dout[i * OC + oc]; + + // Each warp is responsible for 32 * "x128::size" = 256 OCs at BF16 (OC must be a multiple of 256!) + // Block size is 512 threads (16 warps) and we reduce those 16 values into 1 at the end + // blockDim.x is 32 --> single warp being responsible for those 256 OCs + // blockDim.y is 16 --> 16 parallel independent warps processing the same OCs for different BTs + // gridDim.x is OC / 256 --> each block processes 256 OCs + // grimDim.y is max(1, (cuda_num_SMs * cuda_threads_per_SM) / (512 * gridDim.x)); --> fill up the entire GPU! + const int block_size = 512; + const int block_size_x = 32; + const int block_size_y = block_size / block_size_x; // 16 + const int OC_per_warp = block_size_x * x128::size; // 256 at BF16 + + int local_oc = threadIdx.x * x128::size; + int global_oc = blockIdx.x * OC_per_warp + local_oc; + float accumulators[x128::size]; + __shared__ float shared[OC_per_warp]; + + for (int k = 0; k < x128::size; k++) { + accumulators[k] = 0.0f; + } + int thread_id = threadIdx.y * block_size_x + threadIdx.x; + for (int i = thread_id; i < OC_per_warp; i += block_size) { + shared[i] = 0.0f; + } + __syncthreads(); + for (int i = blockIdx.y*block_size_y + threadIdx.y; i < B * T; i += gridDim.y*block_size_y) { + x128 packed_dout = load128(dout + global_oc + i*OC); + for (int k = 0; k < x128::size; k++) { + //printf("%d: %f + %f\n", oc, accumulators[k], (float)packed_dout[k]); + accumulators[k] += (float)packed_dout[k]; + } + //__syncthreads(); // keep block synchronised to maximise memory locality (?) + } + for (int k = 0; k < x128::size; k++) { + atomicAdd(shared + local_oc + k, accumulators[k]); + } + __syncthreads(); + if (threadIdx.y == 0) { + for (int i = threadIdx.x; i < OC_per_warp; i += block_size_x) { + //printf("%d => %f\n", i, shared[i]); + atomicAdd(dbias + i + blockIdx.x*OC_per_warp, shared[i]); + } } - // and atomically add everything together. atomics within one block are conflict-free! - atomicAdd(dbias + oc, sum); } __global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, @@ -1169,9 +1215,6 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl 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; // go through blocks in reverse order, so the slowest block starts first @@ -1181,10 +1224,6 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl datt += idx * T * T; dpreatt += idx * T * T; - if (warp.meta_group_rank() == 0) { - block_acc[warp.thread_rank()] = 0; - } - for(int to = 0; to < T_per_block; ++to) { int t = t0 - to; if(t < 0) return; @@ -1193,15 +1232,13 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl floatX* dpreatt_bth = dpreatt + t * T; float local_sum = 0; - for (int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { + for (int t2 = threadIdx.x; t2 <= t; t2 += BlockSize) { local_sum += (float)att_bth[t2] * (float)datt_bth[t2]; } - 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{}); + local_sum = blockReduce(local_sum, warpReduceSum); - for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) { + for (int t3 = threadIdx.x; 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 = (float)__ldcs(att_bth + t3) * ((float)__ldcs(datt_bth + t3) - local_sum); @@ -1264,7 +1301,7 @@ __device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, i // do the loop in reverse to maximise probability of L2 cache hits // so even small L2s get some hits on the 2nd read of the same thread for (int i = (V+x128::size-1)/x128::size + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) { - x128 packed_x = load128cs(x + i * x128::size); // load and do not keep in cache + x128 packed_x = load128(x + i * x128::size); // try to keep in cache until next read for(int k = 0; k < packed_x.size; ++k) { if (i*x128::size+k >= V) { // bounds checking against real V continue; @@ -1276,14 +1313,6 @@ __device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, i thread_sumval += expf(v - thread_maxval); } } - // two reductions of up to 1024 threads: - // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) - // this results in much cleaner assembly than a multi-warp cg::reduce - __shared__ float shared_maxval[32]; - __shared__ float shared_sumval[32]; - int num_warps = blockDim.x / 32; - int warp_id = threadIdx.x / 32; - int lane_id = threadIdx.x % 32; // Block Max Reduction -> Maths -> Block Sum Reduction float block_maxval = blockReduce(thread_maxval, warpReduceMax); @@ -1299,7 +1328,7 @@ __device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, i __global__ void fused_classifier_kernel3(floatX* logits, floatX* losses, floatX* probs, const floatX* dlosses, const int* targets, int B, int T, int V, int P) { - int idx = blockIdx.x; + int idx = gridDim.x - (blockIdx.x+1); // reverse order for cache hits on matmul data int ix = targets[idx]; // softmax (reading B * T * V, same logits read again below, hopefully still in cache) @@ -1364,7 +1393,7 @@ void encoder_forward(floatX* out, const int block_size = 256; const int N = B * T * C; const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size)); - encoder_forward_kernel3<<>>(out, inp, wte, wpe, B, T, C); + encoder_forward_kernel3<<>>(out, inp, wte, wpe, B, T, C); cudaCheck(cudaGetLastError()); } @@ -1375,7 +1404,7 @@ void encoder_backward(floatX* dwte, floatX* dwpe, const int N = B * T * C; const int block_size = 256; const int grid_size = CEIL_DIV(N, block_size); - encoder_backward_kernel<<>>(dwte, dwpe, dout, inp, B, T, C); + encoder_backward_kernel<<>>(dwte, dwpe, dout, inp, B, T, C); cudaCheck(cudaGetLastError()); } @@ -1386,7 +1415,7 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, const int block_size = 512; const int N = B * T; const int grid_size = CEIL_DIV(N * 32, block_size); - layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); + layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); cudaCheck(cudaGetLastError()); } @@ -1494,7 +1523,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, v = qkvr + 2 * B * T * C; int total_threads = B * NH * T * HS; int num_blocks = CEIL_DIV(total_threads, block_size); - permute_kernel<<>>(q, k, v, inp, B, T, NH, HS); + permute_kernel<<>>(q, k, v, inp, B, T, NH, HS); cudaCheck(cudaGetLastError()); // IMPORTANT: alpha/beta are FP32 for CUBLAS_COMPUTE_32F even if FP16 inputs/outputs @@ -1522,7 +1551,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, // multiply all elements of preatt elementwise by scale float scale = 1.0 / sqrtf(HS); int grid_size = CEIL_DIV(B * NH * T * 32, softmax_block_size); - softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); + softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); cudaCheck(cudaGetLastError()); // new approach: first cuBLAS another batched matmul @@ -1543,7 +1572,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, // now unpermute // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side num_blocks = CEIL_DIV(B * T * C, block_size); - unpermute_kernel<<>>(vaccum, out, B, T, NH, HS); + unpermute_kernel<<>>(vaccum, out, B, T, NH, HS); cudaCheck(cudaGetLastError()); } @@ -1551,7 +1580,7 @@ void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) { NVTX_RANGE_FN(); const int block_size = 256; const int grid_size = CEIL_DIV(N, block_size * x128::size); - residual_forward_kernel<<>>(out, inp1, inp2, N); + residual_forward_kernel<<>>(out, inp1, inp2, N); cudaCheck(cudaGetLastError()); } @@ -1559,7 +1588,7 @@ void gelu_forward(floatX* out, const floatX* inp, int N) { NVTX_RANGE_FN(); const int block_size = 512; const int grid_size = CEIL_DIV(N, block_size * x128::size); - gelu_forward_kernel2<<>>(out, inp, N); + gelu_forward_kernel2<<>>(out, inp, N); cudaCheck(cudaGetLastError()); } @@ -1567,7 +1596,7 @@ void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const in NVTX_RANGE_FN(); const int block_size = 128; const int grid_size = CEIL_DIV(N, block_size * x128::size); - gelu_backward_kernel<<>>(dinp, inp, dout, N); + gelu_backward_kernel<<>>(dinp, inp, dout, N); cudaCheck(cudaGetLastError()); } @@ -1578,6 +1607,42 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, NVTX_RANGE_FN(); float one = 1.0f; float zero = 0.0f; + + // backward to bias, if given, does a += + if (dbias != NULL) { + // Each warp is responsible for 32 * "x128::size" = 256 OCs at BF16 (OC must be a multiple of 256!) + // Block size is 512 threads (16 warps) and we reduce those 16 values into 1 at the end + // blockDim.x is 32 --> single warp being responsible for those 256 OCs + // blockDim.y is 16 --> 16 parallel independent warps processing the same OCs for different BTs + // gridDim.x is OC / 256 --> each block processes 256 OCs + // grimDim.y is max(1, (cuda_num_SMs * cuda_threads_per_SM) / (512 * gridDim.x)); --> fill up the entire GPU! + const int warp_size = 32; + const int block_size = 512; + const int OC_per_warp = warp_size * x128::size; // 256 at BF16 + const int block_size_x = 32; + const int block_size_y = block_size / block_size_x; // 16 + const int grid_size_x = OC / OC_per_warp; // e.g. 3 horizontal blocks for 768 OCs at BF16 + const int grid_size_y = max(1, cuda_threads_per_SM * cuda_num_SMs / (block_size * grid_size_x)); // full GPU! + + assert((OC % OC_per_warp) == 0); // there is no bounds checking in the kernel to maximise performance + + // Run the memset on a parallel stream then immediately wait on it (runs in parallel with previous kernel) + cudaMemsetAsync(dbias_buffer, 0, OC * sizeof(float), parallel_streams[0]); + cudaEventRecord(parallel_events[0], parallel_streams[0]); + cudaStreamWaitEvent(main_stream, parallel_events[0], 0); + + matmul_backward_bias_kernel6<<>>(dbias_buffer, dout, B, T, OC); + cudaEventRecord(main_event, main_stream); + cudaStreamWaitEvent(parallel_streams[0], main_event, 0); // cast_and_add_kernel is dependent on this + + // run cast_and_add kernel on separate stream so it can run in parallel with the matmuls + cast_and_add_kernel<<>>(dbias, dbias_buffer, OC); + cudaEventRecord(parallel_events[0], parallel_streams[0]); + cudaCheck(cudaGetLastError()); + } + // backward to input, uses = in the backward pass (set the gradient) cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one, weight, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &zero, @@ -1586,16 +1651,10 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one, inp, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &one, dweight, CUBLAS_LOWP, C, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - // backward to bias, if given, does a += + if (dbias != NULL) { - const int block_size = 128; - const int grid_size_x = CEIL_DIV(OC, block_size); - const int grid_size_y = max(1, cuda_threads_per_SM * cuda_num_SMs / block_size); - cudaMemset(dbias_buffer, 0, OC * sizeof(float)); - matmul_backward_bias_kernel5<<>>(dbias_buffer, dout, B, T, OC); - cudaCheck(cudaGetLastError()); - cast_and_add_kernel<<>>(dbias, dbias_buffer, OC); - cudaCheck(cudaGetLastError()); + // make sure cast_and_add_kernel is done (hopefully finished in parallel with the matmuls) + cudaStreamWaitEvent(main_stream, parallel_events[0], 0); } } @@ -1606,9 +1665,17 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr const int block_size = 1024; const int grid_size = 1 * cuda_num_SMs; size_t shared_mem_size = (2 * C + 1) * sizeof(float); - cudaMemset(scratch, 0, (2 * C + 1) * sizeof(float)); // todo - memset in parallel with previous kernels using streams - layernorm_backward_kernel7<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); + + cudaMemsetAsync(scratch, 0, (2 * C + 1) * sizeof(float), parallel_streams[0]); + cudaEventRecord(parallel_events[0], parallel_streams[0]); + cudaStreamWaitEvent(main_stream, parallel_events[0], 0); // cast_and_add_kernel is dependent on this + + layernorm_backward_kernel7<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); cudaCheck(cudaGetLastError()); + + // create a dependency between this kernel and future uses of the scratch buffer (e.g. in matmul_backward) + cudaEventRecord(main_event, main_stream); + cudaStreamWaitEvent(parallel_streams[0], main_event, 0); // cast_and_add_kernel is dependent on this } // the sequence of transformations in this compound op is: @@ -1641,7 +1708,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da // backward through the unpermute operation int num_blocks = CEIL_DIV(B * T * C, block_size); - unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); + unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); cudaCheck(cudaGetLastError()); // backward into datt @@ -1657,7 +1724,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da // backward into preatt int hs = C / NH; // head size float scale = 1.0f / sqrtf(hs); - softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); + softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); cudaCheck(cudaGetLastError()); // backward into q cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, alpha_ptr, @@ -1669,7 +1736,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da dk, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); // backward into inp num_blocks = CEIL_DIV(B * NH * T * HS, block_size); - permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); + permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); cudaCheck(cudaGetLastError()); } @@ -1682,7 +1749,7 @@ void fused_classifier3(Type* logits, Type* losses, const int block_size = 1024; const int N = B * T; const int grid_size = N; - fused_classifier_kernel3<<>>(logits, losses, (Type*)NULL, dlosses, targets, B, T, V, P); + fused_classifier_kernel3<<>>(logits, losses, (Type*)NULL, dlosses, targets, B, T, V, P); cudaCheck(cudaGetLastError()); } @@ -1794,7 +1861,7 @@ typedef struct { floatX* lnf; // (B, T, C) floatX* lnf_mean; // (B, T) floatX* lnf_rstd; // (B, T) - floatX* losses; // (B, T) + floatX* losses; // (B, T) // todo - no longer used as GPU writes directly to cpu_losses // adding these two compared to the CPU .c code, needed for attention kernel as buffers floatX* qkvr; // (L, B, T, 3*C) // in inference mode, this buffer will store the logits @@ -2062,9 +2129,10 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { } // copy inputs/targets to the model - cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpyAsync(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice, main_stream)); if (targets != NULL) { - cudaCheck(cudaMemcpy(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpyAsync(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice, parallel_streams[0])); + cudaEventRecord(parallel_events[0], parallel_streams[0]); } // forward pass @@ -2140,16 +2208,18 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { // also forward the cross-entropy loss function if we have the targets if (targets != NULL) { NvtxRange classifier_and_loss_range("classifier_and_loss"); + // wait on memcpy of targets (it definitely finished by now, but better safe than sorry) + cudaStreamWaitEvent(main_stream, parallel_events[0], 0); // fused classifier: does the forward pass and first part of the backward pass // we're passing dlosses = NULL, which will default them to 1.0f/(B*T), i.e. uniform loss - fused_classifier3(acts.output, acts.losses, (floatX*)NULL, model->targets, B, T, V, Vp); - // for convenience also evaluate the mean loss (TODO re-think this compute+sync point) - // move the (B,T) losses to CPU - cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(floatX), cudaMemcpyDeviceToHost)); - float mean_loss = 0.0f; - for (int i=0; icpu_losses[i]); } - mean_loss /= B*T; - model->mean_loss = mean_loss; + fused_classifier3(acts.output, model->cpu_losses, (floatX*)NULL, model->targets, B, T, V, Vp); + + // the GPU now writes the losses directly to the CPU buffer allocated with cudaMallocHost() + // we accumulate cpu_losses at the end of gpt2_backward() waiting on this event + cudaEventRecord(loss_event, main_stream); + + // reset mean_loss here so gpt2_backward() knows we have targets + model->mean_loss = 0.0f; } else { // if we don't have targets, we don't have loss model->mean_loss = -1.0f; @@ -2158,8 +2228,12 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { void gpt2_zero_grad(GPT2 *model) { NVTX_RANGE_FN(); - if (model->grads_acts_memory != NULL) { cudaCheck(cudaMemset(model->grads_acts_memory, 0, model->num_grad_acts * sizeof(floatX))); } - if (model->grads_memory != NULL) { cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX))); } + if (model->grads_memory != NULL) { + cudaCheck(cudaMemsetAsync(model->grads_memory, 0, model->num_parameters * sizeof(floatX), parallel_streams[0])); + } + // Allow this to run in parallel with forward pass, but create a dependency with everything after (backwards pass) + cudaEventRecord(parallel_events[0], parallel_streams[0]); + cudaStreamWaitEvent(main_stream, parallel_events[0], 0); } void gpt2_backward(GPT2 *model) { @@ -2204,6 +2278,15 @@ void gpt2_backward(GPT2 *model) { ActivationTensors acts = model->acts; GradActTensors grads_acts = model->grads_acts; + // reset residual stream gradients (put here to work with gradient accumulation) + cudaCheck(cudaMemsetAsync(model->grads_acts.residual3, 0, B * T * C * sizeof(floatX), parallel_streams[0])); + // allow the memset to run in parallel with the forward pass, but create a dependency with everything after + cudaEventRecord(parallel_events[0], parallel_streams[0]); + cudaStreamWaitEvent(main_stream, parallel_events[0], 0); + // synchronise the other way around as well to avoid read-after-write hazards for the scratch buffer + cudaEventRecord(main_event, main_stream); + cudaStreamWaitEvent(parallel_streams[0], main_event, 0); // cast_and_add_kernel is dependent on this + // re-use the output buffer of the forward pass as a scratchpad during backward pass float* scratchF = (float*)acts.output; @@ -2292,6 +2375,11 @@ void gpt2_backward(GPT2 *model) { layernorm_backward(dresidual, dl_ln1w, dl_ln1b, scratchF, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); } encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C); + + // accumulate the loss, this was calculated at the end of gpt2_forward() + cudaCheck(cudaEventSynchronize(loss_event)); // hopefully finished long ago + for (int i=0; imean_loss += (float)(model->cpu_losses[i]); } + model->mean_loss /= B*T; } // Compute a mean of a single CPU value across all GPU processes. No-op when multi-GPU is disabled. @@ -2318,7 +2406,7 @@ void gpt2_multi_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) { model->num_parameters, ncclFloatX, ncclAvg, multi_gpu_config->nccl_comm, - // use 0 for default stream (all other computations use this stream) + // use 0 for default stream (always implicitly synchronised) /*stream=*/0)); #endif } @@ -2338,7 +2426,7 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo if (model->use_master_weights == 1) { // allocate one more buffer to keep the master copy of weights as float, and copy the weights over cudaCheck(cudaMalloc((void**)&model->master_weights, model->num_parameters * sizeof(float))); - copy_and_cast_kernel<<num_parameters, 512), 512, 0, stream>>>(model->master_weights, (floatX*)model->params_memory, model->num_parameters); + copy_and_cast_kernel<<num_parameters, 512), 512, 0, main_stream>>>(model->master_weights, (floatX*)model->params_memory, model->num_parameters); cudaCheck(cudaGetLastError()); printf0("allocated %zu MiB for master copy of params\n", (model->num_parameters * sizeof(float)) >> 20); } @@ -2349,7 +2437,7 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo float beta1_correction = 1.0f - powf(beta1, t); float beta2_correction = 1.0f - powf(beta2, t); unsigned int seed = random_u32(&model->rng_state); - adamw_kernel3<<>>((floatX*)model->params_memory, model->master_weights, + adamw_kernel3<<>>((floatX*)model->params_memory, model->master_weights, (floatX*)model->grads_memory, model->m_memory, model->v_memory, model->num_parameters, learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, seed); @@ -2593,9 +2681,12 @@ int main(int argc, char *argv[]) { cuda_arch_major = deviceProp.major; cuda_arch_minor = deviceProp.minor; - cudaCheck(cudaStreamCreate(&stream)); + cudaCheck(cudaStreamCreate(&main_stream)); + cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming); + cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming); for (int i = 0; i < num_parallel_streams; i++) { cudaCheck(cudaStreamCreate(¶llel_streams[i])); + cudaEventCreateWithFlags(¶llel_events[i], cudaEventDisableTiming); } // set up cuBLAS and cuBLASLt @@ -2768,6 +2859,7 @@ int main(int argc, char *argv[]) { } gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1); + // todo - move or double-buffer all of this timing logic to avoid idling the GPU at this point! cudaEventRecord(end); float time_elapsed_ms; cudaCheck(cudaEventSynchronize(end)); // wait for the end event to finish to get correct timings