aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix build error

Browse source
This commit is contained in:
Mainak-learner 2025-10-01 15:40:55 -05:00
parent 6839568e36
commit cd9366e9aa
2 changed files with 28 additions and 30 deletions
Download patch file Download diff file Expand all files Collapse all files

44
llmc/layernorm.cuh
View file

@ -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<x128*>(params);
x128* s_bias = reinterpret_cast<x128*>(params) + (C / x128::size);
x128* s_res = reinterpret_cast<x128*>(params) + ((2 + threadIdx.y) * C / x128::size);
x128* s_res = reinterpret_cast<x128*>(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<<<grid_size, dim3(WARP_SIZE, block_y), smem, stream>>>(residual, normed,
mean, rstd, inp1, inp2,
weight, bias, N, C);
fused_residual_forward_kernel5<<<grid_size, dim3(WARP_SIZE, block_y), smem, stream>>>(
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());
}

14
train_gpt2.cu
View file

@ -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);