/* Kernels for layernorm backward pass. Compile example: nvcc -O3 --use_fast_math -lcublas -lcublasLt 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 #include #include #include #include #include #define ENABLE_BF16 #include "common.h" // ---------------------------------------------------------------------------- // CPU code reference void layernorm_forward_cpu(float* out, float* mean, float* rstd, const float* inp, const float* weight, const float* bias, int B, int T, int C) { // reference: https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html // both inp and out are (B,T,C) of the activations // mean and rstd are (B,T) buffers, to be used later in backward pass // at each position (b,t) of the input, the C-dimensional vector // of activations gets normalized, then scaled and shifted float eps = 1e-5f; for (int b = 0; b < B; b++) { for (int t = 0; t < T; t++) { // seek to the input position inp[b,t,:] const float* x = inp + b * T * C + t * C; // calculate the mean float m = 0.0f; for (int i = 0; i < C; i++) { m += x[i]; } m = m/C; // calculate the variance (without any bias correction) float v = 0.0f; for (int i = 0; i < C; i++) { float xshift = x[i] - m; v += xshift * xshift; } v = v/C; // calculate the rstd (reciprocal standard deviation) float s = 1.0f / sqrtf(v + eps); // seek to the output position in out[b,t,:] float* out_bt = out + b * T * C + t * C; for (int i = 0; i < C; i++) { float n = (s * (x[i] - m)); // normalize float o = n * weight[i] + bias[i]; // scale and shift out_bt[i] = o; // write } // cache the mean and rstd for the backward pass later mean[b * T + t] = m; rstd[b * T + t] = s; } } } void layernorm_backward_cpu(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) { for (int b = 0; b < B; b++) { for (int t = 0; t < 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]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = 0; i < C; i++) { float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = weight[i] * dout_bt[i]; dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } dnorm_mean = dnorm_mean / C; dnorm_norm_mean = dnorm_norm_mean / C; // now iterate again and accumulate all the gradients for (int i = 0; i < C; i++) { float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = weight[i] * dout_bt[i]; // gradient contribution to bias dbias[i] += dout_bt[i]; // gradient contribution to weight dweight[i] += norm_bti * dout_bt[i]; // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] += dval; } } } } // ---------------------------------------------------------------------------- // GPU kernels // GPU helper functions for atomicAdd on smaller than 32-bit types #ifdef ENABLE_BF16 __device__ void atomicAddX(__nv_bfloat16* addr, __nv_bfloat16 val) { uintptr_t ptr_val = reinterpret_cast(addr); __nv_bfloat162* ptr_bf16 = reinterpret_cast<__nv_bfloat162*>(ptr_val & ~uintptr_t(0x3)); // Prepare the value to add, setting the other half to zero __nv_bfloat162 add_val = (ptr_val & 0x3) ? __halves2bfloat162(__ushort_as_bfloat16(0), val) : __halves2bfloat162(val, __ushort_as_bfloat16(0)); atomicAdd(ptr_bf16, add_val); } #endif #ifdef ENABLE_FP16 __device__ void atomicAddX(half* addr, half val) { uintptr_t ptr_val = reinterpret_cast(addr); half2* ptr_fp16 = reinterpret_cast(ptr_val & ~uintptr_t(0x3)); // Prepare the value to add, setting the other half to zero half2 add_val = (ptr_val & 0x3) ? __halves2half2(__ushort_as_half(0), val) : __halves2half2(val, __ushort_as_half(0)); atomicAdd(ptr_fp16, add_val); } #endif __device__ void atomicAddX(float* addr, float val) { atomicAdd(addr, val); } // super naive kernel that just parallelizes over B,T and loops over C __global__ void layernorm_backward_kernel1(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) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= B*T) return; int b = idx / T; int t = idx % 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]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = 0; i < C; i++) { float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = weight[i] * dout_bt[i]; dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } dnorm_mean = dnorm_mean / C; dnorm_norm_mean = dnorm_norm_mean / C; // now iterate again and accumulate all the gradients for (int i = 0; i < C; i++) { 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]); // gradient contribution to weight atomicAdd(&dweight[i], norm_bti * dout_bt[i]); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] += dval; } } // uses shared memory instead for the reduces template __global__ void layernorm_backward_kernel2(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, float* dweight_tmp, float* dbias_tmp) { 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; } // 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; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)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; for (int i = warp.thread_rank(); i < C; i += warp.size()) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; 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; // now iterate again and accumulate all the gradients for (int i = warp.thread_rank(); i < C; i += warp.size()) { 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 atomicAdd(&dbias_shared[i], (float)dout_bt[i]); // gradient contribution to weight atomicAdd(&dweight_shared[i], norm_bti * (float)dout_bt[i]); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 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_tmp[i], dbias_shared[i]); atomicAdd(&dweight_tmp[i], dweight_shared[i]); } } template __global__ void copy_to_dweight_dbias(int C, Tparams* dbias, Tparams* dweight, float* dbias_tmp, float* dweight_tmp) { for (int i = threadIdx.x + blockDim.x * blockIdx.x; i < C; i += blockDim.x * gridDim.x) { dbias[i] = (Tparams)dbias_tmp[i]; dweight[i] = (Tparams)dweight_tmp[i]; } } // kernel2 is 1 threadblock for all Cs on 32 BTs (assuming threadblock size of 1024 threads = 32 warps) // To minimise the amount of atomicAdds, we will aim for 1 threadblock per SM, processing (total BTs / threadblocks) BTs template __global__ void layernorm_backward_kernel3(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) { 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 base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // 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 4 for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } __syncthreads(); int warps_in_grid = gridDim.x * warp.meta_group_size(); for (int idx = base_idx; idx < B * T; idx += warps_in_grid) { 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; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warp.thread_rank(); i < C; i += warp.size()) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; 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; // now iterate again and accumulate all the gradients for (int i = warp.thread_rank(); i < C; i += warp.size()) { float dout_i = (float)__ldcs(&dout_bt[i]); float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * dout_i; // gradient contribution to bias atomicAdd(&dbias_shared[i], dout_i); // gradient contribution to weight atomicAdd(&dweight_shared[i], norm_bti * dout_i); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval); } } __syncthreads(); for(int i = threadIdx.x; i < C; i+= blockDim.x) { atomicAddX(&dbias[i], (Tparams)dbias_shared[i]); atomicAddX(&dweight[i], (Tparams)dweight_shared[i]); } } // atomicCAS version of kernel3 template __global__ void layernorm_backward_kernel4(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) { 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 base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // 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 4 for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } __syncthreads(); int warps_in_grid = gridDim.x * warp.meta_group_size(); for (int idx = base_idx; idx < B * T; idx += warps_in_grid) { 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; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warp.thread_rank(); i < C; i += warp.size()) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; 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; // now iterate again and accumulate all the gradients for (int i = warp.thread_rank(); i < C; i += warp.size()) { float dout_i = (float)__ldcs(&dout_bt[i]); float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * dout_i; // gradient contribution to bias atomicAdd(&dbias_shared[i], dout_i); // gradient contribution to weight atomicAdd(&dweight_shared[i], norm_bti * dout_i); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval); } } __syncthreads(); __nv_bfloat162* dbiasVec2 = reinterpret_cast<__nv_bfloat162*>(dbias); __nv_bfloat162* dweightVec2 = reinterpret_cast<__nv_bfloat162*>(dweight); // write to global memory for(int i = threadIdx.x; i < C/2; i+= blockDim.x) { __nv_bfloat162 add_dbias = __halves2bfloat162((__nv_bfloat16)dbias_shared[i*2], (__nv_bfloat16)dbias_shared[i*2+1]); __nv_bfloat162 add_dweight = __halves2bfloat162((__nv_bfloat16)dweight_shared[i*2], (__nv_bfloat16)dweight_shared[i*2+1]); // Get the current value from L2 cache __nv_bfloat162 current_dbias = __ldcg(&dbiasVec2[i]); __nv_bfloat162 current_dweight = __ldcg(&dweightVec2[i]); // Add the two values __nv_bfloat162 new_dbias = add_dbias + current_dbias; __nv_bfloat162 new_dweight = add_dweight + current_dweight; // Write the result back to L2 cache using 32-bit integer atomic compare and exchange unsigned int current_dbias32b = *reinterpret_cast(¤t_dbias); unsigned int current_dweight32b = *reinterpret_cast(¤t_dweight); unsigned int new_dbias32b = *reinterpret_cast(&new_dbias); unsigned int new_dweight32b = *reinterpret_cast(&new_dweight); unsigned int old_dbias32b = atomicCAS((unsigned int*)&dbiasVec2[i], current_dbias32b, new_dbias32b); unsigned int old_dweight32b = atomicCAS((unsigned int*)&dweightVec2[i], current_dweight32b, new_dweight32b); // If the value has changed between read and atomic, we need to try again while (old_dbias32b != current_dbias32b) { current_dbias32b = old_dbias32b; new_dbias = *reinterpret_cast<__nv_bfloat162*>(¤t_dbias32b) + add_dbias; new_dbias32b = *reinterpret_cast(&new_dbias); old_dbias32b = atomicCAS((unsigned int*)&dbiasVec2[i], current_dbias32b, new_dbias32b); } while (old_dweight32b != current_dweight32b) { current_dweight32b = old_dweight32b; new_dweight = *reinterpret_cast<__nv_bfloat162*>(¤t_dweight32b) + add_dweight; new_dweight32b = *reinterpret_cast(&new_dweight); old_dweight32b = atomicCAS((unsigned int*)&dweightVec2[i], current_dweight32b, new_dweight32b); } } } // FP32 scratchpad per threadgroup, zero atomics except atomicAdd on unsigned int for the flag (based on kernel3) template __global__ void layernorm_backward_kernel5(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C) { extern __shared__ float shared[]; // size = 2 * C + 1 namespace cg = cooperative_groups; cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); int base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // 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 4 for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } unsigned int *tmp_flag = (unsigned int*)(shared + C*2); __syncthreads(); int warps_in_grid = gridDim.x * warp.meta_group_size(); for (int idx = base_idx; idx < B * T; idx += warps_in_grid) { 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; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warp.thread_rank(); i < C; i += warp.size()) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; 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; // now iterate again and accumulate all the gradients for (int i = warp.thread_rank(); i < C; i += warp.size()) { float dout_i = (float)__ldcs(&dout_bt[i]); float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * dout_i; // gradient contribution to bias atomicAdd(&dbias_shared[i], dout_i); // gradient contribution to weight atomicAdd(&dweight_shared[i], norm_bti * dout_i); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval); } } __syncthreads(); float* scratch_dbias = scratch; float* scratch_dweight = scratch + C * gridDim.x; unsigned int* scratchFlag = (unsigned int*)(scratch + (2 * C * gridDim.x)); for(int i = threadIdx.x; i < C; i+= blockDim.x) { scratch_dbias[i + C*blockIdx.x] = dbias_shared[i]; scratch_dweight[i + C*blockIdx.x] = dweight_shared[i]; } __threadfence(); __syncthreads(); if (threadIdx.x == 0) { *tmp_flag = atomicAdd(scratchFlag, 1); } __syncthreads(); if (*tmp_flag == gridDim.x-1) { // last block to finish, accumulate the scratchpad for (int i = threadIdx.x; i < C; i += blockDim.x) { float dbias_sum = 0.0f; float dweight_sum = 0.0f; #pragma unroll 8 for (int j = 0; j < gridDim.x; j++) { dbias_sum += scratch_dbias[i + j*C]; dweight_sum += scratch_dweight[i + j*C]; } dbias[i] = (Tparams)((float)dbias[i] + dbias_sum); dweight[i] = (Tparams)((float)dweight[i] + dweight_sum); } } } // single FP32 scratchpad shared by all the threadblocks (based on kernels 3 & 5) template __global__ void layernorm_backward_kernel6(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C) { extern __shared__ float shared[]; // size = 2 * C + 1 namespace cg = cooperative_groups; cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); int base_idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // 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 4 for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } unsigned int *tmp_flag = (unsigned int*)(shared + C*2); __syncthreads(); int warps_in_grid = gridDim.x * warp.meta_group_size(); for (int idx = base_idx; idx < B * T; idx += warps_in_grid) { 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; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warp.thread_rank(); i < C; i += warp.size()) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; 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; // now iterate again and accumulate all the gradients for (int i = warp.thread_rank(); i < C; i += warp.size()) { float dout_i = (float)__ldcs(&dout_bt[i]); float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * dout_i; // gradient contribution to bias atomicAdd(&dbias_shared[i], dout_i); // gradient contribution to weight atomicAdd(&dweight_shared[i], norm_bti * dout_i); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval); } } // Accumulate into a FP32 scratchpad // BF16 atomics are potentially much slower... and this is more precise! __syncthreads(); float* scratch_dbias = scratch; float* scratch_dweight = scratch + C; unsigned int* scratchFlag = (unsigned int*)(scratch + (2 * C)); for(int i = threadIdx.x; i < C; i+= blockDim.x) { atomicAdd(&scratch_dbias[i], dbias_shared[i]); atomicAdd(&scratch_dweight[i], dweight_shared[i]); } __syncthreads(); if (threadIdx.x == 0) { *tmp_flag = atomicAdd(scratchFlag, 1); } __syncthreads(); if (*tmp_flag == gridDim.x-1) { for(int i = threadIdx.x; i < C; i+= blockDim.x) { // todo - potentially do stochastic rounding here as well dbias[i] = (Tparams)scratch_dbias[i]; dweight[i] = (Tparams)scratch_dweight[i]; } } } // Same as kernel 6 but without cooperative groups or templates __global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { extern __shared__ float shared[]; // size = 2 * C + 1 int warpId = threadIdx.x / warpSize; // warp index within a block int warpsInBlock = blockDim.x / warpSize; int base_idx = blockIdx.x * warpsInBlock + warpId; int warpThreadIdx = threadIdx.x % warpSize; // Thread index within the warp int warps_in_grid = gridDim.x * warpsInBlock; // 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 4 for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } unsigned int *tmp_flag = (unsigned int*)(shared + C*2); __syncthreads(); for (int idx = base_idx; idx < B * T; idx += warps_in_grid) { int b = idx / T; int t = idx % T; const floatX* dout_bt = dout + b * T * C + t * C; const floatX* inp_bt = inp + b * T * C + t * C; floatX* dinp_bt = dinp + b * T * C + t * C; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warpThreadIdx; i < C; i += warpSize) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } dnorm_mean = warpReduceSum(dnorm_mean); dnorm_norm_mean = warpReduceSum(dnorm_norm_mean); dnorm_mean = dnorm_mean / C; dnorm_norm_mean = dnorm_norm_mean / C; // now iterate again and accumulate all the gradients for (int i = warpThreadIdx; i < C; i += warpSize) { float dout_i = (float)__ldcs(&dout_bt[i]); float norm_bti = ((float)__ldcs(&inp_bt[i]) - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * dout_i; // gradient contribution to bias atomicAdd(&dbias_shared[i], dout_i); // gradient contribution to weight atomicAdd(&dweight_shared[i], norm_bti * dout_i); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] = (floatX)((float)dinp_bt[i] + dval); } } // Accumulate into a FP32 scratchpad // BF16 atomics are potentially much slower... and this is more precise! __syncthreads(); float* scratch_dbias = scratch; float* scratch_dweight = scratch + C; unsigned int* scratchFlag = (unsigned int*)(scratch + (2 * C)); for(int i = threadIdx.x; i < C; i+= blockDim.x) { atomicAdd(&scratch_dbias[i], dbias_shared[i]); atomicAdd(&scratch_dweight[i], dweight_shared[i]); } __syncthreads(); if (threadIdx.x == 0) { *tmp_flag = atomicAdd(scratchFlag, 1); } __syncthreads(); if (*tmp_flag == gridDim.x-1) { for(int i = threadIdx.x; i < C; i+= blockDim.x) { // todo - potentially do stochastic rounding here as well dbias[i] = (floatX)scratch_dbias[i]; dweight[i] = (floatX)scratch_dweight[i]; } } } __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS) layernorm_backward_kernel8(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { extern __shared__ float shared[]; // size = 2 * C + 1 int warpId = threadIdx.x / warpSize; // warp index within a block int warpsInBlock = blockDim.x / warpSize; //number of warps in block int baseIdx = blockIdx.x * warpsInBlock + warpId; int warpThreadIdx = threadIdx.x % warpSize; // Thread index within the warp int warpsInGrid = gridDim.x * warpsInBlock; int C_per_iteration = warpSize * x128::size; int iterations_C = C / C_per_iteration; // 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 for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } unsigned int *tmp_flag = (unsigned int*)(shared + C*2); __syncthreads(); for (int idx = baseIdx; idx < B * T; idx += warpsInGrid) { int b = idx / T; int t = idx % T; const floatX* dout_bt = dout + b * T * C + t * C; const floatX* inp_bt = inp + b * T * C + t * C; floatX* dinp_bt = dinp + b * T * C + t * C; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warpThreadIdx * x128::size; i < C; i += warpSize * x128::size) { x128 dout128_i = load128(dout_bt + i); x128 inp128_i = load128(inp_bt + i); x128 weight128_i = load128(weight + i); for (int k = 0; k < x128::size; k++) { float norm_bti = ((float)inp128_i[k] - mean_bt) * rstd_bt; float dnorm_i = (float)weight128_i[k] * (float)dout128_i[k]; dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } } dnorm_mean = warpReduceSum(dnorm_mean) / C; dnorm_norm_mean = warpReduceSum(dnorm_norm_mean) / C; // now iterate again and accumulate all the gradients // unfortunately we cannot use the same index for x128 arrays and shared memory // as atomics can only be 32-bit rather than 128-bit (at least pre-SM90/Hopper) // so this would result in an 8-way bank conflict, and kill performance // so instead, we use a shared memory friendly index, and reorder before the final write for (int i = 0; i < iterations_C; i++) { int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration); int shared_index = warpThreadIdx + (i * C_per_iteration); x128 dout128 = load128cs(dout_bt + global_index); x128 inp128 = load128cs(inp_bt + global_index); x128 dinp128 = load128(dinp_bt + global_index); x128 weight128 = load128(weight + global_index); for (int x = 0; x < x128::size; x++) { float dout_i = (float)dout128[x]; float norm_bti = ((float)inp128[x] - mean_bt) * rstd_bt; float dnorm_i = (float)weight128[x] * dout_i; // gradient contribution to bias (using shared memory friendly index) atomicAdd(&dbias_shared[shared_index + x*warpSize], dout_i); // gradient contribution to weight (using shared memory friendly index) atomicAdd(&dweight_shared[shared_index + x*warpSize], norm_bti * dout_i); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp128[x] = (floatX)((float)dinp128[x] + dval); } // cache in L2 as this is read by the next kernel, but bypass L1 to minimise thrashing store128cg(dinp_bt + global_index, dinp128); } } // Accumulate into a FP32 scratchpad // BF16 atomics are potentially much slower... and this is more precise! // todo - could potentially avoid the extra copy if floatX is FP32, fairly negligible though __syncthreads(); float* scratch_dbias = scratch; float* scratch_dweight = scratch + C; unsigned int* scratchFlag = (unsigned int*)(scratch + (2 * C)); for(int i = threadIdx.x; i < C; i+= blockDim.x) { // global atomics in the same "shared memory banking friendly" order atomicAdd(&scratch_dbias[i], dbias_shared[i]); atomicAdd(&scratch_dweight[i], dweight_shared[i]); } __syncthreads(); if (threadIdx.x == 0) { *tmp_flag = atomicInc(scratchFlag, gridDim.x); } __syncthreads(); if (*tmp_flag == gridDim.x-1) { for (int i = warpId; i < iterations_C; i += warpsInBlock) { // reorder from atomic/shared memory-friendly index to real global memory index // and convert from float/FP32 to floatX/BF16 for the final write int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration); int shared_index = warpThreadIdx + (i * C_per_iteration); x128 dbias128 = load128(dbias + global_index); x128 dweight128 = load128(dweight + global_index); for (int x = 0; x < x128::size; x++) { float s_db = scratch_dbias[shared_index + x*warpSize]; float s_dw = scratch_dweight[shared_index + x*warpSize]; dbias128[x] = (floatX)(s_db + (float)dbias128[x]); dweight128[x] = (floatX)(s_dw + (float)dweight128[x]); } store128(dbias + global_index, dbias128); store128(dweight + global_index, dweight128); } } } __global__ void layernorm_backward_kernel9(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { if(C % (32 * x128::size) != 0) { if(threadIdx.x == 0 && blockIdx.x == 0) { printf("Number of channels is not a multiple of 32 * x128::size"); } __trap(); // prefer to crash here than run into a deadlock later on } constexpr int WARP_SIZE = 32; int BLOCK_SIZE = blockDim.x; int warpsInBlock = BLOCK_SIZE / WARP_SIZE; //number of warps in block extern __shared__ float shared[]; // size = 2 * C + 1 int warpId = threadIdx.x / WARP_SIZE; // warp index within a block int baseIdx = blockIdx.x * warpsInBlock + warpId; int warpThreadIdx = threadIdx.x % WARP_SIZE; // Thread index within the warp int warpsInGrid = gridDim.x * warpsInBlock; int C_per_iteration = WARP_SIZE * x128::size; int iterations_C = ceil_div(C, C_per_iteration) + 2; // the first half of shared memory is bias, second is weight float* dbias_shared = shared; float* dweight_shared = shared + C; float* dbias_tmp_shared = shared + 2 * C; float* dweight_tmp_shared = shared + 2 * C + BLOCK_SIZE; // init shared memory to zero for(int i = threadIdx.x; i < C; i+= BLOCK_SIZE){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } unsigned int *tmp_flag = (unsigned int*)(shared + 2*C + 2*BLOCK_SIZE); __syncthreads(); for (int idx = baseIdx; idx < B * T; idx += warpsInGrid) { int b = idx / T; int t = idx % T; const floatX* dout_bt = dout + b * T * C + t * C; const floatX* inp_bt = inp + b * T * C + t * C; floatX* dinp_bt = dinp + b * T * C + t * C; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warpThreadIdx * x128::size; i < C; i += WARP_SIZE * x128::size) { x128 dout128_i = load128(dout_bt + i); x128 inp128_i = load128(inp_bt + i); x128 weight128_i = load128(weight + i); for (int k = 0; k < x128::size; k++) { float norm_bti = ((float)inp128_i[k] - mean_bt) * rstd_bt; float dnorm_i = (float)weight128_i[k] * (float)dout128_i[k]; dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * norm_bti; } } dnorm_mean = warpReduceSum(dnorm_mean) / C; dnorm_norm_mean = warpReduceSum(dnorm_norm_mean) / C; // now iterate again and accumulate all the gradients // unfortunately we cannot use the same index for x128 arrays and shared memory // as atomics can only be 32-bit rather than 128-bit (at least pre-SM90/Hopper) // so this would result in an 8-way bank conflict, and kill performance // so instead, we use a shared memory friendly index, and reorder before the final write for (int i = 0; i < iterations_C; i++) { int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration); int shared_index = warpThreadIdx + (i * C_per_iteration); if (global_index >= C) { break; } x128 dout128 = load128cs(dout_bt + global_index); x128 inp128 = load128cs(inp_bt + global_index); x128 dinp128 = load128(dinp_bt + global_index); x128 weight128 = load128(weight + global_index); for (int x = 0; x < x128::size; x++) { float dout_i = (float)dout128[x]; float norm_bti = ((float)inp128[x] - mean_bt) * rstd_bt; float dnorm_i = (float)weight128[x] * dout_i; // sum up the gradients for bias and weight across the entire block // this is basically a reduction (but only inter-warp, not intra-warp) // doing it this way allows us to avoid using atomics while using many warps if (warpId != 0) { dbias_tmp_shared[threadIdx.x] = dout_i; dweight_tmp_shared[threadIdx.x] = norm_bti * dout_i; } __syncthreads(); if (warpId == 0) { float dbias_tmp = dout_i; float dweight_tmp = norm_bti * dout_i; for (int j = 1; j < warpsInBlock; j++) { dbias_tmp += dbias_tmp_shared[threadIdx.x + j * WARP_SIZE]; dweight_tmp += dweight_tmp_shared[threadIdx.x + j * WARP_SIZE]; } // gradient contribution to bias (using shared memory friendly index) dbias_shared[shared_index + x*WARP_SIZE] += dbias_tmp; // gradient contribution to weight (using shared memory friendly index) dweight_shared[shared_index + x*WARP_SIZE] += dweight_tmp; } __syncthreads(); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp128[x] = (floatX)((float)dinp128[x] + dval); } // cache in L2 as this is read by the next kernel, but bypass L1 to minimise thrashing store128cg(dinp_bt + global_index, dinp128); } } __syncthreads(); // Each block writes its partial sum to global memory // The last block to finish becomes responsible for summing up all the partial sums // This is done by atomically incrementing a flag (cleared to 0 before launching the kernel) unsigned int* scratchFlag = (unsigned int*)(scratch); // Increment scratch pointer by a full cacheline so that everything remains cacheline aligned scratch += 32; float* scratch_dbias = scratch; float* scratch_dweight = scratch + C; for(int i = threadIdx.x; i < C; i+= BLOCK_SIZE) { // Write to global memory in the same "shared memory banking friendly" order scratch_dbias[i + 2*C*blockIdx.x] = dbias_shared[i]; scratch_dweight[i + 2*C*blockIdx.x] = dweight_shared[i]; } __syncthreads(); if (threadIdx.x == 0) { *tmp_flag = atomicInc(scratchFlag, gridDim.x); } __syncthreads(); if (*tmp_flag == gridDim.x-1) { // Reduction of the partial sums by the final block // todo - there isn't enough parallelism even inside that single SM... // ==> so could maybe split into another kernel with YET ANOTHER level of reduction?! for(int i = threadIdx.x * f128::size; i < C; i+= BLOCK_SIZE * f128::size) { f128 dbias_accum = f128::zeros(); f128 dweight_accum = f128::zeros(); for (int read_block_idx = 0; read_block_idx < gridDim.x; read_block_idx++) { int offset = i + 2*C*read_block_idx; f128 dbias128 = load128(scratch_dbias + offset); f128 dweight128 = load128(scratch_dweight + offset); for(int k = 0; k < f128::size; k++) { dbias_accum[k] += dbias128[k]; dweight_accum[k] += dweight128[k]; } } store128(dbias_shared + i, dbias_accum); store128(dweight_shared + i, dweight_accum); } __syncthreads(); // reorder from atomic/shared memory-friendly index to real global memory index // and convert from float/FP32 to floatX/BF16 for the final write // this is separate also because it cannot use as many warps as the above (f128 vs x128) // todo - if we split this code into another kernel, we could maybe do it at the same time? for (int i = warpId; i < iterations_C; i += warpsInBlock) { int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration); int shared_index = warpThreadIdx + (i * C_per_iteration); if (global_index >= C) { break; } x128 dbias128 = load128(dbias + global_index); x128 dweight128 = load128(dweight + global_index); for (int x = 0; x < x128::size; x++) { float s_db = dbias_shared[shared_index + x*WARP_SIZE]; float s_dw = dweight_shared[shared_index + x*WARP_SIZE]; dbias128[x] = (floatX)(s_db + (float)dbias128[x]); dweight128[x] = (floatX)(s_dw + (float)dweight128[x]); } store128(dbias + global_index, dbias128); store128(dweight + global_index, dweight128); } } } // similar to kernel 9, but uses vectors to access shared memory, which also avoids the bank conflict problems, // and makes use require fewer barriers, at the cost of increased shared memory consumption. // warning: this kernel is _extremely_ close to getting register spills, so many "optimizations" turn out to be unhelpful // or need to be implemented in a very specific way. __global__ void __launch_bounds__(512, 2) layernorm_backward_kernel10(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { constexpr int WARP_SIZE = 32; int BLOCK_SIZE = blockDim.x; int warpsInBlock = BLOCK_SIZE / WARP_SIZE; //number of warps in block extern __shared__ float shared[]; // size = 2 * C + 1 int warpId = threadIdx.x / WARP_SIZE; // warp index within a block int baseIdx = blockIdx.x * warpsInBlock + warpId; int warpThreadIdx = threadIdx.x % WARP_SIZE; // Thread index within the warp int warpsInGrid = gridDim.x * warpsInBlock; int C_per_iteration = WARP_SIZE * x128::size; int iterations_C = ceil_div(C, C_per_iteration); // + 2; // the first half of shared memory is bias, second is weight size_t rounded_C = ceil_div(C, (32 * x128::size)) * (32 * x128::size); float* dbias_shared = shared; float* dweight_shared = shared + rounded_C; // warp zero doesn't actually write to the _tmp_shared memory locations, so we don't need to reserve memory // the obvious solution is to change the addressing below to use (threadId.x-32) as offset, but that causes // register spills, so instead we mess with the base pointer here, which doesn't increase register usage. float* dbias_tmp_shared = shared + 2 * rounded_C - WARP_SIZE * f128::size; float* dweight_tmp_shared = shared + 2 * rounded_C + f128::size * BLOCK_SIZE - 2 * WARP_SIZE * f128::size; // init shared memory to zero for(int i = threadIdx.x * f128::size; i < rounded_C; i += BLOCK_SIZE * f128::size) { store128(dbias_shared + i, f128::zeros()); store128(dweight_shared + i, f128::zeros()); } __syncthreads(); for (int bt = baseIdx; bt < B * T; bt += warpsInGrid) { const floatX* dout_bt = dout + bt * C; const floatX* inp_bt = inp +bt * C; floatX* dinp_bt = dinp + bt * C; // first: two reduce operations float dnorm_mean = 0.0f; float dnorm_norm_mean = 0.0f; for (int i = warpThreadIdx * x128::size; i < C; i += WARP_SIZE * x128::size) { x128 dout128_i = load128(dout_bt + i); x128 inp128_i = load128(inp_bt + i); x128 weight128_i = load128(weight + i); for (int k = 0; k < x128::size; k++) { float dnorm_i = (float)weight128_i[k] * (float)dout128_i[k]; dnorm_mean += dnorm_i; dnorm_norm_mean += dnorm_i * (float)inp128_i[k]; } } const float mean_bt = (float)mean[bt]; const float rstd_bt = (float)rstd[bt]; dnorm_mean = warpReduceSum(dnorm_mean) / C; dnorm_norm_mean = warpReduceSum(dnorm_norm_mean) / C * rstd_bt - dnorm_mean * mean_bt * rstd_bt; for (int c = 0; c < iterations_C; c++) { int global_index = (warpThreadIdx * x128::size) + (c * C_per_iteration); x128 dout128 = x128::zeros(); x128 inp128 = x128::zeros(); x128 dinp128 = x128::zeros(); x128 weight128 = x128::zeros(); if(global_index < C) { dout128 = load128cs(dout_bt + global_index); inp128 = load128cs(inp_bt + global_index); dinp128 = load128(dinp_bt + global_index); weight128 = load128(weight + global_index); } for(int o = 0; o < x128::size / f128::size; ++o) { f128 dbias_f; f128 dweight_f; for(int i = 0; i < f128::size; ++i) { int x = o * f128::size + i; float dout_i = (float)dout128[x]; float norm_bti = ((float)inp128[x] - mean_bt) * rstd_bt; dbias_f[i] = dout_i; dweight_f[i] = norm_bti * dout_i; float dval = 0.0f; dval += (float) weight128[x] * (float)dout128[x]; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp128[x] = (floatX) ((float) dinp128[x] + dval); } if (warpId != 0) { store128(dbias_tmp_shared + threadIdx.x * f128::size, dbias_f); // this seems to generate a 64-bit store, instead of 128-bit. // however, forcing 128-bit (e.g., using inline ptx), results in register // spilling and much worse performance, so we'll keep it like this for now // but ideally, we could reduce the register pressure a little. store128(dweight_tmp_shared + threadIdx.x * f128::size, dweight_f); } __syncthreads(); if (warpId == 0) { for (int j = 1; j < warpsInBlock; j++) { f128 dbias_tmp = load128(dbias_tmp_shared + f128::size * (threadIdx.x + j * WARP_SIZE)); f128 dweight_tmp = load128(dweight_tmp_shared + f128::size * (threadIdx.x + j * WARP_SIZE)); for(int i = 0; i < f128::size; ++i) { dbias_f[i] += dbias_tmp[i]; dweight_f[i] += dweight_tmp[i]; } } } __syncthreads(); if (warpId == 0) { f128 db_old = load128(dbias_shared + global_index + f128::size * o); f128 dw_old = load128(dweight_shared + global_index + f128::size * o); for(int i = 0; i < f128::size; ++i) { dbias_f[i] += db_old[i]; dweight_f[i] += dw_old[i]; } store128(dbias_shared + global_index + f128::size * o, dbias_f); store128(dweight_shared + global_index + f128::size * o, dweight_f); } } if(global_index < C) { // cache in L2 as this is read by the next kernel, but bypass L1 to minimise thrashing store128cg(dinp_bt + global_index, dinp128); } } } __syncthreads(); // Each block writes its partial sum to global memory // The last block to finish becomes responsible for summing up all the partial sums // This is done by atomically incrementing a flag (cleared to 0 before launching the kernel) unsigned int* scratchFlag = (unsigned int*)(scratch); // Increment scratch pointer by a full cacheline so that everything remains cacheline aligned scratch += 32; float* scratch_dbias = scratch; float* scratch_dweight = scratch + C; for(int i = threadIdx.x * f128::size; i < C; i += BLOCK_SIZE * f128::size) { // Write to global memory in the same "shared memory banking friendly" order store128(scratch_dbias + i + 2*C*blockIdx.x, load128(dbias_shared + i)); store128(scratch_dweight + i + 2*C*blockIdx.x, load128(dweight_shared + i)); } __syncthreads(); // that portion of shared memory is no longer used, so we can repurpose it for the scratch flag. unsigned int *tmp_flag = (unsigned int*)(shared + 2*rounded_C); if (threadIdx.x == 0) { *tmp_flag = atomicInc(scratchFlag, gridDim.x); } __syncthreads(); if (*tmp_flag == gridDim.x-1) { // Reduction of the partial sums by the final block // todo - there isn't enough parallelism even inside that single SM... // ==> so could maybe split into another kernel with YET ANOTHER level of reduction?! for(int i = threadIdx.x * f128::size; i < C; i += BLOCK_SIZE * f128::size) { f128 dbias_accum = f128::zeros(); f128 dweight_accum = f128::zeros(); for (int read_block_idx = 0; read_block_idx < gridDim.x; read_block_idx++) { int offset = i + 2*C*read_block_idx; f128 dbias128 = load128(scratch_dbias + offset); f128 dweight128 = load128(scratch_dweight + offset); for(int k = 0; k < f128::size; k++) { dbias_accum[k] += dbias128[k]; dweight_accum[k] += dweight128[k]; } } store128(dbias_shared + i, dbias_accum); store128(dweight_shared + i, dweight_accum); } __syncthreads(); // convert from float/FP32 to floatX/BF16 for the final write // this is separate because it cannot use as many warps as the above (f128 vs x128) // todo - if we split this code into another kernel, we could maybe do it at the same time? for (int c = warpId; c < iterations_C; c += warpsInBlock) { int global_index = (warpThreadIdx * x128::size) + (c * C_per_iteration); if (global_index >= C) { break; } x128 dbias128 = load128(dbias + global_index); x128 dweight128 = load128(dweight + global_index); for(int o = 0; o < x128::size / f128::size; ++o) { f128 s_db = load128(dbias_shared + global_index + o * f128::size); f128 s_dw = load128(dweight_shared + global_index + o * f128::size); for(int i = 0; i < f128::size; ++i) { int x = o * f128::size + i; dbias128[x] = (floatX)(s_db[i] + (float)dbias128[x]); dweight128[x] = (floatX)(s_dw[i] + (float)dweight128[x]); } } store128(dbias + global_index, dbias128); store128(dweight + global_index, dweight128); } } } // ---------------------------------------------------------------------------- // kernel launchers void layernorm_backward1(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, const int block_size) { const int N = B * T; const int grid_size = ceil_div(N, block_size); layernorm_backward_kernel1<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward2(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, int block_size) { const int N = B * T; const int grid_size = ceil_div(32*N, block_size); size_t shared_mem_size = 2 * C * sizeof(float); float* dweight_tmp; float* dbias_tmp; cudaCheck(cudaMalloc(&dweight_tmp, C * sizeof(float))); cudaCheck(cudaMalloc(&dbias_tmp, C * sizeof(float))); cudaMemset(dweight_tmp, 0, C * sizeof(float)); cudaMemset(dbias_tmp, 0, C * sizeof(float)); layernorm_backward_kernel2<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, dweight_tmp, dbias_tmp); copy_to_dweight_dbias<<<1, 512>>>(C, dweight, dbias, dweight_tmp, dbias_tmp); cudaCheck(cudaFree(dweight_tmp)); cudaCheck(cudaFree(dbias_tmp)); } template void layernorm_backward3(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, int block_size) { const int grid_size = (1024/block_size) * cuda_num_SMs; size_t shared_mem_size = 2 * C * sizeof(float); layernorm_backward_kernel3<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward4(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, int block_size) { const int grid_size = (1024/block_size) * cuda_num_SMs; size_t shared_mem_size = 2 * C * sizeof(float); layernorm_backward_kernel4<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward5(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C, int block_size) { const int grid_size = 1 * cuda_num_SMs; // only support 1 block per SM for simplicity, 1024 threads is best anyway size_t shared_mem_size = (2 * C + 1) * sizeof(float); cudaMemset(scratch, 0, (grid_size * 2 * C + 1) * sizeof(float)); layernorm_backward_kernel5<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward6(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C, int block_size) { const int grid_size = (1024/block_size) * cuda_num_SMs; size_t shared_mem_size = (2 * C + 1) * sizeof(float); // Including this as part of the timing until we can parallelise it // It should fully hide the cost and improve kernel perf by >5% if done in parallel using CUDA streams cudaMemset(scratch, 0, (1 + 2 * C) * sizeof(float)); layernorm_backward_kernel6<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward7(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C, int block_size) { const int grid_size = (1024/block_size) * cuda_num_SMs; size_t shared_mem_size = (2 * C + 1) * sizeof(float); // Including this as part of the timing until we can parallelise it // It should fully hide the cost and improve kernel perf by >5% if done in parallel using CUDA streams cudaMemset(scratch, 0, (1 + 2 * C) * sizeof(float)); layernorm_backward_kernel7<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward8(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C, int block_size) { const int grid_size = (1024/block_size) * cuda_num_SMs; size_t shared_mem_size = (2 * C + 1) * sizeof(float); // Including this as part of the timing until we can parallelise it // It should fully hide the cost and improve kernel perf by >5% if done in parallel using CUDA streams cudaMemset(scratch, 0, (1 + 2 * C) * sizeof(float)); layernorm_backward_kernel8<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward9(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C, int block_size) { assert(C % (32 * x128::size) == 0 && "Channels must be divisible by (32 * x128::size)"); const int grid_size = (1024/block_size) * cuda_num_SMs; // todo - heuristics for other GPUs? size_t shared_mem_size = (2 * C + 2 * block_size + 1) * sizeof(float); cudaMemset(scratch, 0, 1 * sizeof(float)); // just need to memset the flag for this version layernorm_backward_kernel9<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); } template void layernorm_backward10(Tdinp* dinp, Tparams* dweight, Tparams* dbias, float* scratch, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C, int block_size) { if(block_size == 1024) { block_size = 512; } //assert(C % (32 * x128::size) == 0 && "Channels must be divisible by (32 * x128::size)"); const int grid_size = (1024/block_size) * cuda_num_SMs; // todo - heuristics for other GPUs? size_t rounded_C = ceil_div(C, (32 * x128::size)) * (32 * x128::size); size_t shared_mem_size = (2 * rounded_C + 2 * (block_size - 32) * f128::size) * sizeof(float); cudaCheck(cudaMemset(scratch, 0, 1 * sizeof(float))); // just need to memset the flag for this version layernorm_backward_kernel10<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); cudaCheck(cudaGetLastError()); } // kernel version dispatch void layernorm_backward(int kernel_num, floatX* dinp, floatX* dweight, floatX* dbias, float* scratch, const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C, const int block_size) { switch (kernel_num) { #if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) case 1: layernorm_backward1(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size); break; #endif case 2: layernorm_backward2(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size); break; case 3: layernorm_backward3(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size); break; #if defined(ENABLE_BF16) case 4: layernorm_backward4(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C, block_size); break; #endif case 5: layernorm_backward5(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); break; case 6: layernorm_backward6(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); break; case 7: layernorm_backward7(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); break; case 8: layernorm_backward8(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); break; case 9: layernorm_backward9(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); break; case 10: layernorm_backward10(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C, block_size); break; default: printf("Invalid kernel number\n"); exit(1); } cudaCheck(cudaGetLastError()); } // ---------------------------------------------------------------------------- int main(int argc, char **argv) { setup_main(); int B = 8; int T = 1024; int C = 1600; // this is the problematic size // first do the forward pass in CPU float* out = (float*)malloc(B * T * C * sizeof(float)); float* mean = (float*)malloc(B * T * sizeof(float)); float* rstd = (float*)malloc(B * T * sizeof(float)); float* inp = make_random_float(B * T * C); float* weight = make_random_float(C); float* bias = make_random_float(C); layernorm_forward_cpu(out, mean, rstd, inp, weight, bias, B, T, C); // now do the backward pass, again on CPU float *dout = make_random_float(B * T * C); float *dinp = make_zeros_float(B * T * C); float *dweight = make_zeros_float(C); float *dbias = make_zeros_float(C); layernorm_backward_cpu(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); // the above calculations act as the reference // now let's do the same on the GPU // read kernel_num from command line int kernel_num = 2; if (argc > 1) { kernel_num = atoi(argv[1]); } printf("Using kernel %d\n", kernel_num); // move all the variables we need for backward pass onto the GPU floatX* d_dinp; floatX* d_dweight; floatX* d_dbias; floatX* d_dout; floatX* d_inp; floatX* d_weight; floatX* d_mean; floatX* d_rstd; float* d_scratch; cudaCheck(cudaMalloc(&d_dinp, B * T * C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_dweight, C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_dbias, C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_dout, B * T * C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_weight, C * sizeof(floatX))); cudaCheck(cudaMalloc(&d_mean, B * T * sizeof(floatX))); cudaCheck(cudaMalloc(&d_rstd, B * T * sizeof(floatX))); cudaCheck(cudaMalloc(&d_scratch, (1024/32) * cuda_num_SMs * (2 * C + 1) * sizeof(float))); // copy over the "inputs" to the backward call cudaCheck(memcpy_convert(d_dout, dout, B * T * C)); cudaCheck(memcpy_convert(d_inp, inp, B * T * C)); cudaCheck(memcpy_convert(d_weight, weight, C)); cudaCheck(memcpy_convert(d_mean, mean, B * T)); cudaCheck(memcpy_convert(d_rstd, rstd, B * T)); // launch the kernel // removed 768 because it doesn't work for kernel9 despite being OK in train_gpt2.cu?! int block_sizes[] = {32, 64, 128, 256, 512, /*768,*/ 1024}; for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { int block_size = block_sizes[j]; // init the "outputs" of the backward call to zeros cudaCheck(cudaMemset(d_dinp, 0, B * T * C * sizeof(floatX))); cudaCheck(cudaMemset(d_dweight, 0, C * sizeof(floatX))); cudaCheck(cudaMemset(d_dbias, 0, C * sizeof(floatX))); layernorm_backward(kernel_num, d_dinp, d_dweight, d_dbias, d_scratch, d_dout, d_inp, d_weight, d_mean, d_rstd, B, T, C, block_size); // check the correctness of the kernel float error_threshold_dinp = sizeof(floatX) == 4 ? 1e-3f : 1e-1f; // allow larger errors for BF16/FP16 float error_threshold_dparams = sizeof(floatX) == 4 ? 1e-3f : 5e-1f; // much, much larger... printf("Checking correctness...\n"); printf("dinp:\n"); validate_result(d_dinp, dinp, "dinp", B * T * C, error_threshold_dinp); printf("dweight:\n"); validate_result(d_dweight, dweight, "dweight", C, error_threshold_dparams); printf("dbias:\n"); validate_result(d_dbias, dbias, "dbias", C, error_threshold_dparams); printf("All results match for block_size=%d.\n\n", block_size); } // now time the kernel for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { int block_size = block_sizes[j]; int repeat_times = 100; float elapsed_time = benchmark_kernel(repeat_times, layernorm_backward, kernel_num, d_dinp, d_dweight, d_dbias, d_scratch, d_dout, d_inp, d_weight, d_mean, d_rstd, B, T, C, block_size); printf("block_size %4d time %.4f ms\n", block_size, elapsed_time); } // cleanups free(out); free(mean); free(rstd); free(inp); free(weight); free(bias); free(dout); free(dinp); free(dweight); free(dbias); cudaCheck(cudaFree(d_dinp)); cudaCheck(cudaFree(d_dweight)); cudaCheck(cudaFree(d_dbias)); cudaCheck(cudaFree(d_dout)); cudaCheck(cudaFree(d_inp)); cudaCheck(cudaFree(d_weight)); cudaCheck(cudaFree(d_mean)); cudaCheck(cudaFree(d_rstd)); cudaCheck(cudaFree(d_scratch)); return 0; }