From dd8c9f5ec9bf268cd0ed562a4f07214d0bfa1199 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 14 May 2024 20:43:51 +0300 Subject: [PATCH 1/2] fix layernorm backward: accumulate weight gradient --- dev/cuda/layernorm_backward.cu | 10 ++++++---- train_gpt2.cu | 10 ++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/dev/cuda/layernorm_backward.cu b/dev/cuda/layernorm_backward.cu index 575e0a9..e221552 100644 --- a/dev/cuda/layernorm_backward.cu +++ b/dev/cuda/layernorm_backward.cu @@ -842,11 +842,13 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS) int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration); int shared_index = warpThreadIdx + (i * C_per_iteration); - x128 dbias128; - x128 dweight128; + x128 dbias128 = load128(dbias + global_index); + x128 dweight128 = load128(dweight + global_index); for (int x = 0; x < x128::size; x++) { - dbias128[x] = (floatX)scratch_dbias[shared_index + x*warpSize]; - dweight128[x] = (floatX)scratch_dweight[shared_index + x*warpSize]; + 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); diff --git a/train_gpt2.cu b/train_gpt2.cu index 3285aa1..469a0d6 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1029,11 +1029,13 @@ __global__ void __launch_bounds__(512, 3) // todo - any warnings on Turing with int global_index = (warpThreadIdx * x128::size) + (i * C_per_iteration); int shared_index = warpThreadIdx + (i * C_per_iteration); - x128 dbias128; - x128 dweight128; + x128 dbias128 = load128(dbias + global_index); + x128 dweight128 = load128(dweight + global_index); for (int x = 0; x < x128::size; x++) { - dbias128[x] = (floatX)scratch_dbias[shared_index + x*warpSize]; - dweight128[x] = (floatX)scratch_dweight[shared_index + x*warpSize]; + 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); From e553e2f084b29bbf7a59006de593ac311d60fc19 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 14 May 2024 20:44:50 +0300 Subject: [PATCH 2/2] update dev/cuda/layernorm_backward and improve `validate_result` to take into account fp epsilon when comparing results --- dev/cuda/common.h | 22 +++++++++- dev/cuda/layernorm_backward.cu | 74 +++++++++++++--------------------- 2 files changed, 48 insertions(+), 48 deletions(-) diff --git a/dev/cuda/common.h b/dev/cuda/common.h index 0c20798..2757c67 100644 --- a/dev/cuda/common.h +++ b/dev/cuda/common.h @@ -3,6 +3,7 @@ #include #include #include +#include template @@ -260,13 +261,25 @@ void validate_result(D* device_result, const T* cpu_reference, const char* name, D* out_gpu = (D*)malloc(num_elements * sizeof(D)); cudaCheck(cudaMemcpy(out_gpu, device_result, num_elements * sizeof(D), cudaMemcpyDeviceToHost)); int nfaults = 0; +#ifndef ENABLE_BF16 + float epsilon = FLT_EPSILON; +#else + float epsilon = 0.079; +#endif for (int i = 0; i < num_elements; i++) { + // Skip masked elements + if(!isfinite(cpu_reference[i])) + continue; + // print the first few comparisons if (i < 5) { printf("%f %f\n", cpu_reference[i], (T)out_gpu[i]); } - // ensure correctness for all elements. We can set an "ignore" mask by writing NaN - if (fabs(cpu_reference[i] - (T)out_gpu[i]) > tolerance && isfinite(cpu_reference[i])) { + // effective tolerance is based on expected rounding error (epsilon), + // plus any specified additional tolerance + float t_eff = tolerance + fabs(cpu_reference[i]) * epsilon; + // ensure correctness for all elements. + if (fabs(cpu_reference[i] - (T)out_gpu[i]) > t_eff) { printf("Mismatch of %s at %d: CPU_ref: %f vs GPU: %f\n", name, i, cpu_reference[i], (T)out_gpu[i]); nfaults ++; if (nfaults >= 10) { @@ -276,6 +289,11 @@ void validate_result(D* device_result, const T* cpu_reference, const char* name, } } + if (nfaults > 0) { + free(out_gpu); + exit(EXIT_FAILURE); + } + // reset the result pointer, so we can chain multiple tests and don't miss trivial errors, // like the kernel not writing to part of the result. // cudaMemset(device_result, 0, num_elements * sizeof(T)); diff --git a/dev/cuda/layernorm_backward.cu b/dev/cuda/layernorm_backward.cu index e221552..90dcb16 100644 --- a/dev/cuda/layernorm_backward.cu +++ b/dev/cuda/layernorm_backward.cu @@ -1014,25 +1014,6 @@ int main(int argc, char **argv) { float *dbias = make_zeros_float(C); layernorm_backward_cpu(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); - // convert all the necessary cpu data to floatX (e.g. bfloat16) - floatX* meanX = (floatX*)malloc(B * T * sizeof(floatX)); - floatX* rstdX = (floatX*)malloc(B * T * sizeof(floatX)); - floatX* doutX = (floatX*)malloc(B * T * C * sizeof(floatX)); - floatX* inpX = (floatX*)malloc(B * T * C * sizeof(floatX)); - floatX* weightX = (floatX*)malloc(C * sizeof(floatX)); - - for (int i = 0; i < B * T; i++) { - meanX[i] = (floatX)mean[i]; - rstdX[i] = (floatX)rstd[i]; - } - for (int i = 0; i < B * T * C; i++) { - doutX[i] = (floatX)dout[i]; - inpX[i] = (floatX)inp[i]; - } - for (int i = 0; i < C; i++) { - weightX[i] = (floatX)weight[i]; - } - // the above calculations act as the reference // now let's do the same on the GPU @@ -1063,33 +1044,39 @@ int main(int argc, char **argv) { cudaCheck(cudaMalloc(&d_rstd, B * T * sizeof(floatX))); cudaCheck(cudaMalloc(&d_scratch, cuda_num_SMs * (2 * C + 1) * sizeof(float))); // copy over the "inputs" to the backward call - cudaCheck(cudaMemcpy(d_dout, doutX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_inp, inpX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_weight, weightX, C * sizeof(floatX), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_mean, meanX, B * T * sizeof(floatX), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_rstd, rstdX, B * T * sizeof(floatX), cudaMemcpyHostToDevice)); - // 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))); + 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 - const int block_size = 256; - 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); + 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))); - // 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 : 20.0f; // 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); + 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 - int block_sizes[] = {32, 64, 128, 256, 512, 1024}; for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { int block_size = block_sizes[j]; int repeat_times = 100; @@ -1110,11 +1097,6 @@ int main(int argc, char **argv) { free(dinp); free(dweight); free(dbias); - free(meanX); - free(rstdX); - free(doutX); - free(inpX); - free(weightX); cudaCheck(cudaFree(d_dinp)); cudaCheck(cudaFree(d_dweight)); cudaCheck(cudaFree(d_dbias));