From cd9366e9aa52e23b0443aa8a211df7820e2901d1 Mon Sep 17 00:00:00 2001 From: Mainak-learner Date: Wed, 1 Oct 2025 15:40:55 -0500 Subject: [PATCH] fix build error --- llmc/layernorm.cuh | 44 +++++++++++++++++++++----------------------- train_gpt2.cu | 14 +++++++------- 2 files changed, 28 insertions(+), 30 deletions(-) diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index b5d033b..7e1fb33 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -144,57 +144,58 @@ __global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, const floatX* weight, int N, int C) { assert(blockDim.x == WARP_SIZE); - - // load weights and biases into shared memory - // do this before we allow any threads to exit! extern __shared__ char* params[]; - // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so - // let's keep everything as x128 x128* s_weight = reinterpret_cast(params); - x128* s_bias = reinterpret_cast(params) + (C / x128::size); - x128* s_res = reinterpret_cast(params) + ((2 + threadIdx.y) * C / x128::size); + x128* s_res = reinterpret_cast(params) + (C / x128::size); // s_bias removed int sidx = (threadIdx.x + WARP_SIZE * threadIdx.y) * x128::size; for(int i = sidx; i < C; i += blockDim.y * WARP_SIZE * x128::size) { s_weight[i/x128::size] = load128(weight + i); - s_bias[i/x128::size] = load128(bias + i); } __syncthreads(); int idx = blockIdx.x * blockDim.y + threadIdx.y; - if(idx > N) return; + if(idx >= N) return; // Use >= to be safe - // adjust pointers to current token residual += C * idx; normed += C * idx; inp1 += C * idx; inp2 += C * idx; const float eps = 1e-5f; - float sum = 0.0f; + float ss = 0.0f; // sum of squares instead of sum for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { const x128 in1 = load128cs(inp1 + c); const x128 in2 = load128cs(inp2 + c); x128 out; for(int k = 0; k < x128::size; ++k) { out[k] = (float)in1[k] + (float)in2[k]; - sum += (float)out[k]; + ss += (float)out[k] * (float)out[k]; // calculate sum of squares } store128cs(residual + c, out); s_res[c / x128::size] = out; } - sum = warpReduceSum(sum); - float m = sum / C; - float v = 0.f; + ss = warpReduceSum(ss) / C; + float s = rsqrtf(ss + eps); for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { const x128 res = s_res[c / x128::size]; + const x128 w = s_weight[c / x128::size]; + x128 out_data; for(int k = 0; k < x128::size; ++k) { - v += ((float)res[k] - m) * ((float)res[k] - m); + float o = (s * (float)res[k]) * (float)w[k]; // RMSNorm logic + out_data[k] = (floatX)o; } + store128cs(normed + c, out_data); } + if(threadIdx.x == 0) { + // mean[idx] = m; // REMOVED + rstd[idx] = s; + } +} + v = warpReduceSum(v) / C; float s = rsqrtf(v + eps); @@ -471,20 +472,17 @@ void fused_residual_forward5(floatX* residual, floatX* normed, float* rstd, const int block_size = 256; int block_y = block_size / WARP_SIZE; const int grid_size = CEIL_DIV(N, block_y); - size_t smem = (2 + block_y) * C * sizeof(floatX); + size_t smem = (1 + block_y) * C * sizeof(floatX); // 1 for weight, not 2 (no bias) - // in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute - // this may fail, in which case we fall back to the smem free implementation. cudaCheck(cudaGetLastError()); auto status = cudaFuncSetAttribute(fused_residual_forward_kernel5, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); cudaCheck(cudaGetLastError()); if(status == cudaSuccess) { - fused_residual_forward_kernel5<<>>(residual, normed, - mean, rstd, inp1, inp2, - weight, bias, N, C); + fused_residual_forward_kernel5<<>>( + residual, normed, rstd, inp1, inp2, weight, N, C); } else { residual_forward(residual, inp1, inp2, N*C, stream); - layernorm_forward(normed, mean, rstd, residual, weight, bias, N, 1, C, stream); + rmsnorm_forward(normed, rstd, residual, weight, N, 1, C, stream); // Call rmsnorm_forward in fallback } cudaCheck(cudaGetLastError()); } diff --git a/train_gpt2.cu b/train_gpt2.cu index 07be534..d380c08 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -680,7 +680,7 @@ void gpt2_forward(GPT2 *model, const int* inputs, size_t B, size_t T) { encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C, main_stream); // encoding goes into residual[0] // first layernorm isn't fused - layernorm_forward((model->recompute < 2) ? acts.ln1 : acts.lnf, acts.ln1_mean, acts.ln1_rstd, acts.encoded, params.ln1w, params.ln1b, B, T, C, main_stream); + rmsnorm_forward((model->recompute < 2) ? acts.ln1 : acts.lnf, acts.ln1_rstd, acts.encoded, params.ln1w, B, T, C, main_stream); for (int l = 0; l < L; l++) { NvtxRange layer_range("Layer", l); @@ -705,7 +705,7 @@ void gpt2_forward(GPT2 *model, const int* inputs, size_t B, size_t T) { floatX* l_atty = acts.atty + l * B * T * C; floatX* l_residual2 = acts.residual2 + l * B * T * C; floatX* l_ln2 = (model->recompute < 2) ? acts.ln2 + l * B * T * C : acts.lnf; - float* l_ln2_mean = acts.ln2_mean + l * B * T; + // float* l_ln2_mean = acts.ln2_mean + l * B * T; float* l_ln2_rstd = acts.ln2_rstd + l * B * T; floatX* l_fch = acts.fch + l * B * T * 4*C; // reuse the same activation buffer at each layer, as we'll re-compute the gelu during backward @@ -737,7 +737,7 @@ void gpt2_forward(GPT2 *model, const int* inputs, size_t B, size_t T) { // OK, fusion across blocks. if(l+1 != L) { floatX* l_ln1 = (model->recompute < 2) ? acts.ln1 + (l + 1) * B * T * C : acts.lnf; - float* l_ln1_mean = acts.ln1_mean + (l + 1) * B * T; + // float* l_ln1_mean = acts.ln1_mean + (l + 1) * B * T; float* l_ln1_rstd = acts.ln1_rstd + (l + 1) * B * T; const floatX* l_ln1w = params.ln1w + (l + 1) * C; // const floatX* l_ln1b = params.ln1b + (l + 1) * C; @@ -871,13 +871,13 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int floatX* dl_fcprojb = grads.fcprojb + l * C; // get the pointers of the activations for this layer floatX* l_ln1 = (model->recompute < 2) ? acts.ln1 + l * B * T * C : acts.lnf; - float* l_ln1_mean = acts.ln1_mean + l * B * T; + // float* l_ln1_mean = acts.ln1_mean + l * B * T; float* l_ln1_rstd = acts.ln1_rstd + l * B * T; floatX* l_qkvr = acts.qkvr + l * B * T * 3*C; floatX* l_atty = acts.atty + l * B * T * C; floatX* l_residual2 = acts.residual2 + l * B * T * C; floatX* l_ln2 = (model->recompute < 2) ? acts.ln2 + l * B * T * C : acts.lnf; - float* l_ln2_mean = acts.ln2_mean + l * B * T; + // float* l_ln2_mean = acts.ln2_mean + l * B * T; float* l_ln2_rstd = acts.ln2_rstd + l * B * T; floatX* l_fch_pre_gelu = acts.fch + l * B * T * 4*C; floatX* l_fch_gelu = (model->recompute < 1) ? acts.fch_gelu + l * B * T * 4*C : acts.fch_gelu; @@ -896,7 +896,7 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C, main_stream, l_fch_pre_gelu, model->gelu_fusion); if(model->recompute >= 2) { // same as gelu above, l_ln1 and l_ln2 are just buffers if recompute >= 2, recompute them here on demand - layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C, main_stream); + rmsnorm_forward(l_ln2, l_ln2_rstd, l_residual2, l_ln2w, B, T, C, main_stream); } matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, scratchF, B, T, C, 4 * C, main_stream); // layernorm backward does += to the dresidual, so it correctly accumulates grad from the MLP block above @@ -914,7 +914,7 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int attention_backward(dl_bt4c, buffer_b, scratchX, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH, main_stream); #endif if(model->recompute >= 2) { - layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C, main_stream); + rmsnorm_forward(l_ln1, l_ln1_rstd, residual, l_ln1w, B, T, C, main_stream); } // QKV parameter gradients matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, scratchF, B, T, C, 3 * C, main_stream);