From ca7caac5c538141021ceb9a7e8a4b1aa44081ead Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 2 Jun 2024 22:29:04 +0300 Subject: [PATCH 1/4] move all kernels into a main_stream in preparation for future parallelism --- llmc/adamw.cuh | 16 +++++ llmc/attention.cuh | 20 +++--- llmc/cudnn_att.cpp | 22 ++++-- llmc/cudnn_att.h | 4 +- llmc/encoder.cuh | 12 ++-- llmc/fused_classifier.cuh | 4 +- llmc/gelu.cuh | 8 +-- llmc/global_norm.cuh | 4 +- llmc/layernorm.cuh | 23 ++++--- llmc/matmul.cuh | 14 ++-- train_gpt2.cu | 138 +++++++++++++++++++++----------------- 11 files changed, 155 insertions(+), 110 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index 89c3512..eb244a6 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -47,3 +47,19 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg // this will be used in the next update if (master_params_memory != NULL) { master_params_memory[idx] = param; } } + +template +void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters, + float learning_rate, float beta1, float beta2, int t, float eps, float weight_decay, + float grad_scale, unsigned int seed, cudaStream_t stream) { + // AdamW update + int block_size = 512; + int num_blocks = CEIL_DIV(num_parameters, block_size); + float beta1_correction = 1.0f - powf(beta1, t); + float beta2_correction = 1.0f - powf(beta2, t); + adamw_kernel3<<>>(params_memory, master_params_memory, grads_memory, + m_memory, v_memory, num_parameters, + learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, + grad_scale, seed); + cudaCheck(cudaGetLastError()); +} \ No newline at end of file diff --git a/llmc/attention.cuh b/llmc/attention.cuh index 746761f..6d4f0e2 100644 --- a/llmc/attention.cuh +++ b/llmc/attention.cuh @@ -190,7 +190,7 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl void attention_forward(floatX* out, floatX* qkvr, floatX* att, floatX* inp, - int B, int T, int C, int NH) { + int B, int T, int C, int NH, cudaStream_t stream) { NVTX_RANGE_FN(); // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer. // Its contents will be overwritten by this function. @@ -209,9 +209,11 @@ 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); floatX* preatt = inp; + cublasCheck(cublasSetStream(cublas_handle, stream)); + cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size)); cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, @@ -223,7 +225,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 * WARP_SIZE, block_size); - softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); + softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); // new approach: first cuBLAS another batched matmul floatX* vaccum = inp; @@ -239,7 +241,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()); } @@ -248,7 +250,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att, void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* datt, floatX* scratch, const floatX* dout, const floatX* qkvr, const floatX* att, - int B, int T, int C, int NH) { + int B, int T, int C, int NH, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 256; int HS = C / NH; // head size @@ -266,8 +268,10 @@ 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); // backward into datt + cublasCheck(cublasSetStream(cublas_handle, stream)); + cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size)); cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, &beta, datt, CUBLAS_LOWP, T, T * T, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT)); @@ -278,7 +282,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); // backward into q cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha, k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta, @@ -289,6 +293,6 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da dk, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_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()); } diff --git a/llmc/cudnn_att.cpp b/llmc/cudnn_att.cpp index cd25732..7217862 100644 --- a/llmc/cudnn_att.cpp +++ b/llmc/cudnn_att.cpp @@ -21,9 +21,16 @@ static_assert(false, "cuDNN is not supported in FP32 mode.") static cudnnHandle_t cudnn_handle; static size_t cudnn_workspace_size = 0; // dynamically allocated as needed (up to 256MiB!) static void* cudnn_workspace = NULL; -#define checkCudnnErr(err) assert((int)err == 0); -static void checkCudnnFE(fe::error_object e, const char *file, int line) { +static void cuDNNCheck(cudnnStatus_t error, const char *file, int line) { + if (error != CUDNN_STATUS_SUCCESS) { + printf("[CUDNN ERROR] at file %s:%d:\n%s\n", file, line, cudnnGetErrorString(error)); + exit(EXIT_FAILURE); + } +}; +#define cuDNNCheck(err) (cuDNNCheck(err, __FILE__, __LINE__)) + +static void checkCudnnFE(const fe::error_object& e, const char *file, int line) { if(!e.is_good()) { printf("[CUDNN ERROR] at file %s:%d:\n%s\n", file, line, e.err_msg.c_str()); exit(EXIT_FAILURE); @@ -211,11 +218,13 @@ auto lookup_cache_or_build_graph_bwd(int B, int NH, int T, int HS) { void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS) float* stats, // output for backward pass: (B, NH, T) floatX* inp, // input: (B, T, 3, NH, HS) QKV - int B, int T, int NH, int C) { + int B, int T, int NH, int C, cudaStream_t stream) { NVTX_RANGE_FN(); int HS = C / NH; // number of features per head bool is_inference_only = (stats == nullptr); + cuDNNCheck(cudnnSetStream(cudnn_handle, stream)); + // Get graph and tensors from cache (or generate it on first use) auto graph = lookup_cache_or_build_graph_fwd(B, NH, T, HS, is_inference_only); @@ -242,7 +251,7 @@ void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS) void attention_backward_cudnn(floatX* dqkvr, // output floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs - int B, int T, int NH, int C) { + int B, int T, int NH, int C, cudaStream_t stream) { NVTX_RANGE_FN(); int HS = C / NH; // number of features per head @@ -269,15 +278,16 @@ void attention_backward_cudnn(floatX* dqkvr, {Attn_scale_UID, &attn_scale_cpu}}; // Execute graph + cuDNNCheck(cudnnSetStream(cudnn_handle, stream)); checkCudnnFE(graph->execute(cudnn_handle, variant_pack, cudnn_workspace)); cudaCheck(cudaGetLastError()); } void create_cudnn() { - checkCudnnErr(cudnnCreate(&cudnn_handle)); + cuDNNCheck(cudnnCreate(&cudnn_handle)); } void destroy_cudnn() { if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } - checkCudnnErr(cudnnDestroy(cudnn_handle)); + cuDNNCheck(cudnnDestroy(cudnn_handle)); } \ No newline at end of file diff --git a/llmc/cudnn_att.h b/llmc/cudnn_att.h index df88005..3184130 100644 --- a/llmc/cudnn_att.h +++ b/llmc/cudnn_att.h @@ -12,10 +12,10 @@ void destroy_cudnn(); void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS) float* stats, // output for backward pass: (B, NH, T) floatX* inp, // input: (B, T, 3, NH, HS) QKV - int B, int T, int NH, int C); + int B, int T, int NH, int C, cudaStream_t stream); void attention_backward_cudnn(floatX* dqkvr, // output floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs - int B, int T, int NH, int C); + int B, int T, int NH, int C, cudaStream_t stream); #endif // CUDNN_ATT_H \ No newline at end of file diff --git a/llmc/encoder.cuh b/llmc/encoder.cuh index c854b68..bb3068e 100644 --- a/llmc/encoder.cuh +++ b/llmc/encoder.cuh @@ -149,12 +149,12 @@ __global__ void wpe_backward_kernel(floatX* dwpe, void encoder_forward(floatX* out, const int* inp, const floatX* wte, const floatX* wpe, - int B, int T, int C) { + int B, int T, int C, cudaStream_t stream) { NVTX_RANGE_FN(); 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()); } @@ -162,7 +162,7 @@ void encoder_forward(floatX* out, void encoder_backward(floatX* dwte, floatX* dwpe, floatX* scratch, // gpu outputs & scratch int* workload_indices, int4* bucket_info, // cpu scratch buffers const floatX* dout, const int* inp, const int* inputs_cpu, // cpu/gpu inputs - int B, int T, int C, unsigned int seed) { + int B, int T, int C, unsigned int seed, cudaStream_t stream) { NVTX_RANGE_FN(); // Launch wpe kernel first (so it runs on the GPU in parallel with the CPU pre-processing for wte) @@ -217,11 +217,11 @@ void encoder_backward(floatX* dwte, floatX* dwpe, floatX* scratch, // gpu output // todo - could use CUDA events (even without streams) to avoid CPU/GPU synchronisation completely int4* d_bucket_info = (int4*)scratch; int* d_workload_indices = (int*)(scratch + B*T*num_c_groups * sizeof(int4)); - cudaCheck(cudaMemcpyAsync(d_bucket_info, bucket_info, num_buckets * sizeof(int4), cudaMemcpyHostToDevice)); - cudaCheck(cudaMemcpy(d_workload_indices, workload_indices, total_items * sizeof(int), cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpyAsync(d_bucket_info, bucket_info, num_buckets * sizeof(int4), cudaMemcpyHostToDevice, stream)); + cudaCheck(cudaMemcpyAsync(d_workload_indices, workload_indices, total_items * sizeof(int), cudaMemcpyHostToDevice, stream)); // Launch wte kernel // todo - profile block sizes on more content (depends on number of buckets and on GPU?) - wte_backward_kernel<256><<>>(dwte, d_bucket_info, d_workload_indices, dout, inp, seed, B, T, C); + wte_backward_kernel<256><<>>(dwte, d_bucket_info, d_workload_indices, dout, inp, seed, B, T, C); cudaCheck(cudaGetLastError()); } diff --git a/llmc/fused_classifier.cuh b/llmc/fused_classifier.cuh index 4a00b00..b13fa35 100644 --- a/llmc/fused_classifier.cuh +++ b/llmc/fused_classifier.cuh @@ -133,11 +133,11 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS) template void fused_classifier(Type* logits, Type* losses, const float dloss, const int* targets, - int B, int T, int V, int P) { + int B, int T, int V, int P, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 1024; const int N = B * T; const int grid_size = N; - fused_classifier_kernel5<<>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P); + fused_classifier_kernel5<<>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P); cudaCheck(cudaGetLastError()); } diff --git a/llmc/gelu.cuh b/llmc/gelu.cuh index d77a8ff..ce9aa12 100644 --- a/llmc/gelu.cuh +++ b/llmc/gelu.cuh @@ -47,20 +47,20 @@ __global__ void gelu_backward_inplace_kernel(floatX* d_in_out, const floatX* inp // ---------------------------------------------------------------------------- // kernel launchers -void gelu_forward(floatX* out, const floatX* inp, int N) { +void gelu_forward(floatX* out, const floatX* inp, int N, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 512; assert(N % block_size == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); - gelu_forward_kernel2<<>>(out, inp); + gelu_forward_kernel2<<>>(out, inp); cudaCheck(cudaGetLastError()); } -void gelu_backward_inplace(floatX* d_in_out, const floatX* inp, const int N) { +void gelu_backward_inplace(floatX* d_in_out, const floatX* inp, const int N, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 128; assert(N % block_size == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); - gelu_backward_inplace_kernel<<>>(d_in_out, inp); + gelu_backward_inplace_kernel<<>>(d_in_out, inp); cudaCheck(cudaGetLastError()); } diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index 50b8606..69bdd66 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -34,7 +34,7 @@ __global__ void global_norm_squared_kernel(float* out, const T* data, size_t cou // kernel launcher template -void global_norm_squared(float* out, const T* values, size_t count) { +void global_norm_squared(float* out, const T* values, size_t count, cudaStream_t stream) { const int block_size = 512; // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. // having one block less than possible is a tiny performance hit, having @@ -45,7 +45,7 @@ void global_norm_squared(float* out, const T* values, size_t count) { assert(grid_size > 0); // gives a better error than letting the call below fail // initialize out with zero cudaCheck(cudaMemset(out, 0, sizeof(float))); - global_norm_squared_kernel<<>>(out, values, count); + global_norm_squared_kernel<<>>(out, values, count); cudaCheck(cudaGetLastError()); } diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index f599ef3..9965cce 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -356,28 +356,28 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, floatX* inp, const floatX* weight, const floatX* bias, - int B, int T, int C) { + int B, int T, int C, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 512; const int N = B * T; const int grid_size = CEIL_DIV(N * WARP_SIZE, 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()); } -void residual_forward(floatX* out, const floatX* inp1, const floatX* inp2, int N) { +void residual_forward(floatX* out, const floatX* inp1, const floatX* inp2, int N, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 256; assert(N % block_size == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); - residual_forward_kernel<<>>(out, inp1, inp2); + residual_forward_kernel<<>>(out, inp1, inp2); cudaCheck(cudaGetLastError()); } void fused_residual_forward5(floatX* residual, floatX* normed, floatX* mean, floatX* rstd, const floatX* inp1, const floatX* inp2, const floatX* weight, const floatX* bias, - int N, int C) { + int N, int C, cudaStream_t stream) { const int block_size = 256; int block_y = block_size / WARP_SIZE; const int grid_size = CEIL_DIV(N, block_y); @@ -389,18 +389,19 @@ void fused_residual_forward5(floatX* residual, floatX* normed, floatX* mean, flo auto status = cudaFuncSetAttribute(fused_residual_forward_kernel5, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); cudaGetLastError(); if(status == cudaSuccess) { - fused_residual_forward_kernel5<<>>(residual, normed, mean, rstd, inp1, inp2, - weight, bias, N, C); + fused_residual_forward_kernel5<<>>(residual, normed, + mean, rstd, inp1, inp2, + weight, bias, N, C); } else { - residual_forward(residual, inp1, inp2, N*C); - layernorm_forward(normed, mean, rstd, residual, weight, bias, N, 1, C); + residual_forward(residual, inp1, inp2, N*C, stream); + layernorm_forward(normed, mean, rstd, residual, weight, bias, N, 1, C, stream); } cudaCheck(cudaGetLastError()); } void layernorm_backward(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) { + int B, int T, int C, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 512; const int blocks_per_sm = 2; // supported on every architecture and less cache thrashing than 3 @@ -409,6 +410,6 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr size_t shared_mem_size = (2 * rounded_C + 2 * (block_size - 32) * f128::size) * sizeof(float); cudaCheck(cudaMemset(scratch, 0, 1 * sizeof(float))); // only need to reset the flag to 0 - layernorm_backward_kernel10<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); + layernorm_backward_kernel10<<>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C); cudaCheck(cudaGetLastError()); } diff --git a/llmc/matmul.cuh b/llmc/matmul.cuh index e53e612..bc97e37 100644 --- a/llmc/matmul.cuh +++ b/llmc/matmul.cuh @@ -104,7 +104,7 @@ __global__ void reduce_add_sum_kernel(floatX* dst, const float* src, size_t n, s // https://docs.nvidia.com/cuda/cublas/#cublasltmatmul void matmul_forward_cublaslt(floatX* out, floatX* inp, floatX* weight, floatX* bias, - int B, int T, int C, int OC) { + int B, int T, int C, int OC, cudaStream_t stream) { NVTX_RANGE_FN(); int has_bias = (bias != NULL); @@ -161,7 +161,7 @@ void matmul_forward_cublaslt(floatX* out, cublasCheck(cublasLtMatmul(cublaslt_handle, operationDesc, &alpha, weight, weightLayout, inp, inputLayout, &beta, out, outputLayout, out, outputLayout, &heuristic.algo, - cublaslt_workspace, cublaslt_workspace_size, 0)); + cublaslt_workspace, cublaslt_workspace_size, stream)); // cleanups cublasCheck(cublasLtMatmulPreferenceDestroy(preference)); @@ -175,7 +175,7 @@ void matmul_forward_cublaslt(floatX* out, void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, floatX* dout, floatX* inp, floatX* weight, float* dbias_buffer, - int B, int T, int C, int OC) { + int B, int T, int C, int OC, cudaStream_t stream) { NVTX_RANGE_FN(); float one = 1.0f, zero = 0.0f; @@ -194,18 +194,20 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, // If we have enough OC that we don't need cross-block reductions, we can skip the bias_buffer accumulation // and write results directly to the output. if(grid_size_y == 1) { - matmul_backward_bias_kernel9<<>>(dbias, dout, B, T, OC, std::bool_constant{}); + matmul_backward_bias_kernel9<<>>(dbias, dout, B, T, OC, std::bool_constant{}); cudaCheck(cudaGetLastError()); } else { // kernel 9 overwrites temp buffer, so no need to memset - matmul_backward_bias_kernel9<<>>(dbias_buffer, dout, B, T, OC, std::bool_constant{}); + matmul_backward_bias_kernel9<<>>(dbias_buffer, dout, B, T, OC, std::bool_constant{}); cudaCheck(cudaGetLastError()); - reduce_add_sum_kernel<<>>(dbias, dbias_buffer, OC, grid_size_y); + reduce_add_sum_kernel<<>>(dbias, dbias_buffer, OC, grid_size_y); cudaCheck(cudaGetLastError()); } } // backward to input, uses = in the backward pass (set the gradient) + cublasCheck(cublasSetStream(cublas_handle, stream)); + cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size)); cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one, weight, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &zero, dinp, CUBLAS_LOWP, C, cublas_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); diff --git a/train_gpt2.cu b/train_gpt2.cu index f9eeac3..c0731fe 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -511,6 +511,7 @@ typedef struct { // todo - if other functions need cpu scratch buffers in the future, reuse as generic scratch? int* workload_indices; // encoder_backward, B*T*num_c_groups (int) int4* bucket_info; // encoder_backward, B*T*num_c_groups (int4) - size for worst case + cudaStream_t main_stream; } GPT2; void gpt2_init_common(GPT2 *model) { @@ -541,6 +542,7 @@ void gpt2_init_common(GPT2 *model) { model->rng_state = 13371337; // used in stochastic rounding model->use_master_weights = 1; // safe default: do keep master weights in fp32 model->recompute = 1; // good default: recompute gelu but not layernorm + cudaStreamCreate(&model->main_stream); } void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { @@ -633,6 +635,8 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { fcloseCheck(model_file); gpt2_init_common(model); + // only return from this function once we are certain the params are ready on the GPU + cudaCheck(cudaDeviceSynchronize()); } void gpt2_build_from_random(GPT2 *model, int depth) { @@ -724,7 +728,8 @@ void gpt2_build_from_random(GPT2 *model, int depth) { gpt2_init_common(model); } -void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, int grad_accum_steps=1) { +void gpt2_forward(GPT2 *model, const int* inputs, const int* targets, size_t B, size_t T, int grad_accum_steps=1) { + // right now, this function is fully synchronous with the host NVTX_RANGE_FN(); // targets are optional and could be NULL // in this function we must be careful and use size_t instead of int, otherwise @@ -737,11 +742,11 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in } // convenience parameters - size_t V = model->config.vocab_size; - size_t Vp = model->config.padded_vocab_size; - size_t L = model->config.num_layers; - size_t NH = model->config.num_heads; - size_t C = model->config.channels; + const size_t V = model->config.vocab_size; + const size_t Vp = model->config.padded_vocab_size; + const size_t L = model->config.num_layers; + const size_t NH = model->config.num_heads; + const size_t C = model->config.channels; // validate inputs, all indices must be in the range [0, V) for(int i = 0; i < B * T; i++) { @@ -780,19 +785,20 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in } // copy inputs/targets to the model - // todo - inputs is copied on default stream so this synchronises CPU/GPU for now - cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice)); + cudaCheck(cudaMemcpyAsync(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice, model->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, model->main_stream)); } + cudaStream_t main_stream = model->main_stream; + // forward pass ParameterTensors params = model->params; // for brevity ActivationTensors acts = model->acts; - encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0] + 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(acts.ln1, acts.ln1_mean, acts.ln1_rstd, acts.encoded, params.ln1w, params.ln1b, B, T, C); + layernorm_forward(acts.ln1, acts.ln1_mean, acts.ln1_rstd, acts.encoded, params.ln1w, params.ln1b, B, T, C, main_stream); for (int l = 0; l < L; l++) { NvtxRange layer_range("Layer", l); @@ -830,22 +836,22 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in // now do the forward pass #ifdef ENABLE_CUDNN float* l_att = (float*)acts.att + l * B * NH * T; // cuDNN needs a smaller FP32 tensor - matmul_forward_cublaslt(l_qkvr, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C); - attention_forward_cudnn(l_atty, (float*)l_att, l_qkvr, B, T, NH, C); + matmul_forward_cublaslt(l_qkvr, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C, main_stream); + attention_forward_cudnn(l_atty, (float*)l_att, l_qkvr, B, T, NH, C, main_stream); #else floatX* l_att = acts.att + l * B * NH * T * T; // these are only needed as scratchpads for the forward pass, but // need not be stored for backward floatX* scratch = (floatX*)acts.output; - matmul_forward_cublaslt(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C); - attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH); + matmul_forward_cublaslt(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C, main_stream); + attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH, main_stream); #endif - matmul_forward_cublaslt(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C); - fused_residual_forward5(l_residual2, l_ln2, l_ln2_mean, l_ln2_rstd, residual, l_attproj, l_ln2w, l_ln2b, B*T, C); - matmul_forward_cublaslt(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C); - gelu_forward(l_fch_gelu, l_fch, B*T*4*C); - matmul_forward_cublaslt(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C); + matmul_forward_cublaslt(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C, main_stream); + fused_residual_forward5(l_residual2, l_ln2, l_ln2_mean, l_ln2_rstd, residual, l_attproj, l_ln2w, l_ln2b, B*T, C, main_stream); + matmul_forward_cublaslt(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C, main_stream); + gelu_forward(l_fch_gelu, l_fch, B*T*4*C, main_stream); + matmul_forward_cublaslt(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C, main_stream); // OK, fusion across blocks. if(l+1 != L) { @@ -855,22 +861,22 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in const floatX* l_ln1w = params.ln1w + (l + 1) * C; const floatX* l_ln1b = params.ln1b + (l + 1) * C; fused_residual_forward5(l_residual3, l_ln1, l_ln1_mean, l_ln1_rstd, l_residual2, l_fcproj, l_ln1w, l_ln1b, - B * T, C); + B * T, C, main_stream); } else { fused_residual_forward5(l_residual3, acts.lnf, acts.lnf_mean, acts.lnf_rstd, l_residual2, l_fcproj, params.lnfw, params.lnfb, - B * T, C); + B * T, C, main_stream); } } - matmul_forward_cublaslt(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp); + matmul_forward_cublaslt(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp, main_stream); // also forward the cross-entropy loss function if we have the targets if (targets != NULL) { NvtxRange classifier_and_loss_range("classifier_and_loss"); // fused classifier: does the forward pass and first part of the backward pass const float dloss = 1.0f / (B * T * grad_accum_steps); // results in the uniform average loss over all elements - fused_classifier(acts.output, acts.losses, dloss, model->targets, B, T, V, Vp); + fused_classifier(acts.output, acts.losses, dloss, model->targets, B, T, V, Vp, main_stream); // for convenience also evaluate the mean loss (TODO re-think this compute+sync point) cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(floatX), cudaMemcpyDeviceToHost)); float mean_loss = 0.0f; @@ -885,13 +891,15 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in // if we don't have targets, we don't have loss model->mean_loss = -1.0f; } + cudaCheck(cudaDeviceSynchronize()); } void gpt2_zero_grad(GPT2 *model) { NVTX_RANGE_FN(); if (model->grads_memory != NULL) { - cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX))); + cudaCheck(cudaMemsetAsync(model->grads_memory, 0, model->num_parameters * sizeof(floatX), model->main_stream)); } + cudaCheck(cudaDeviceSynchronize()); } void gpt2_backward(GPT2 *model, int* inputs) { @@ -928,12 +936,12 @@ void gpt2_backward(GPT2 *model, int* inputs) { } // convenience shortcuts, size_t instead of int so that pointer arithmetics don't overflow - size_t B = model->batch_size; - size_t T = model->seq_len; - size_t Vp = model->config.padded_vocab_size; - size_t L = model->config.num_layers; - size_t NH = model->config.num_heads; - size_t C = model->config.channels; + const size_t B = model->batch_size; + const size_t T = model->seq_len; + const size_t Vp = model->config.padded_vocab_size; + const size_t L = model->config.num_layers; + const size_t NH = model->config.num_heads; + const size_t C = model->config.channels; // backward pass: go in the reverse order of the forward pass, and call backward() functions ParameterTensors params = model->params; // for brevity @@ -942,22 +950,24 @@ void gpt2_backward(GPT2 *model, int* inputs) { GradActTensors grads_acts = model->grads_acts; // reset residual stream gradients (put here to work with gradient accumulation) - cudaCheck(cudaMemset(model->grads_acts.residual3, 0, B * T * C * sizeof(floatX))); + cudaCheck(cudaMemsetAsync(model->grads_acts.residual3, 0, B * T * C * sizeof(floatX), model->main_stream)); // re-use the output buffer of the forward pass as a scratchpad during backward pass float* scratchF = (float*)acts.output; floatX* scratchX = (floatX*)acts.output; + cudaStream_t main_stream = model->main_stream; + // we kick off the chain rule by filling in dlosses with 1.0f/(B*T) // this was done in the fused classifier kernel as last step of forward pass // technically that is a small, inline backward() pass of calculating // total, final loss as the mean over all losses over all (B,T) positions in the batch // next: backward the classifier matmul - matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp); + matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp, main_stream); // backward the final layernorm floatX* residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 floatX* dresidual = (floatX*)grads_acts.residual3; // the main buffer holding the gradient in the backward pass - layernorm_backward(dresidual, grads.lnfw, grads.lnfb, scratchF, grads_acts.bt4c, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C); + layernorm_backward(dresidual, grads.lnfw, grads.lnfb, scratchF, grads_acts.bt4c, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C, main_stream); // from this point on, we no longer need the values stored in the last residual, so we can reuse that memory as generic // scratch for backward computations @@ -1013,40 +1023,42 @@ void gpt2_backward(GPT2 *model, int* inputs) { if(model->recompute >= 1) { // recompute >= 1 means we recompute gelu. in this case, // l_fch_gelu is just a buffer, so re-compute the gelu from l_fch here - gelu_forward(l_fch_gelu, l_fch, B*T*4*C); + gelu_forward(l_fch_gelu, l_fch, B*T*4*C, main_stream); } - matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C); - gelu_backward_inplace(dl_bt4c, l_fch, B*T*4*C); + matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, scratchF, B, T, 4*C, C, main_stream); + gelu_backward_inplace(dl_bt4c, l_fch, B*T*4*C, main_stream); 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); + layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, 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); + 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 - layernorm_backward(dresidual, dl_ln2w, dl_ln2b, scratchF, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C); - matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, scratchF, B, T, C, C); + layernorm_backward(dresidual, dl_ln2w, dl_ln2b, scratchF, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C, main_stream); + matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, scratchF, B, T, C, C, main_stream); #ifdef ENABLE_CUDNN float* l_att = (float*)acts.att + l * B * NH * T; // cuDNN needs a smaller FP32 tensor - attention_backward_cudnn(dl_bt4c, dl_btc, l_qkvr, l_atty, (float*)l_att, B, T, NH, C); + attention_backward_cudnn(dl_bt4c, dl_btc, l_qkvr, l_atty, (float*)l_att, B, T, NH, C, main_stream); #else floatX* l_att = acts.att + l * B * NH * T * T; // we need B x T x (4)C buffers. l_atty and l_fch aren't needed anymore at this point, so reuse their memory floatX* buffer_a = l_atty; floatX* buffer_b = l_fch; // this is B x T x 4C, so even larger than what we need floatX* dl_preatt = (floatX*)grads_acts.preatt; // dedicated scratchpad allocation - attention_backward(dl_bt4c, buffer_b, dl_preatt, scratchX, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH); + attention_backward(dl_bt4c, buffer_b, dl_preatt, 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); + layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, 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); + matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, scratchF, B, T, C, 3 * C, main_stream); // layernorm backward does += to dresidual, so it correctly accumulates gradient for the Attention block above - layernorm_backward(dresidual, dl_ln1w, dl_ln1b, scratchF, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); + layernorm_backward(dresidual, dl_ln1w, dl_ln1b, scratchF, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C, main_stream); } encoder_backward(grads.wte, grads.wpe, scratchX, model->workload_indices, model->bucket_info, - dresidual, model->inputs, inputs, B, T, C, random_u32(&model->rng_state)); + dresidual, model->inputs, inputs, B, T, C, random_u32(&model->rng_state), main_stream); + + cudaCheck(cudaDeviceSynchronize()); } // Compute sum of a single CPU value across all GPU processes. No-op when multi-GPU is disabled. @@ -1074,16 +1086,17 @@ void gpt2_multi_gpu_grad_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) { ncclCheck(ncclAllReduce(model->grads_memory, model->grads_memory, model->num_parameters, ncclFloatX, ncclAvg, - multi_gpu_config->nccl_comm, 0)); + multi_gpu_config->nccl_comm, model->main_stream)); } else if (multi_gpu_config->zero_stage == 1) { // ZERO-1: Get average gradient for local shard floatX* local_grads_memory = (floatX*) model->grads_memory + multi_gpu_config->shard_offset; ncclCheck(ncclReduceScatter(model->grads_memory, local_grads_memory, multi_gpu_config->shard_num_parameters, ncclFloatX, ncclAvg, - multi_gpu_config->nccl_comm, 0)); + multi_gpu_config->nccl_comm, model->main_stream)); } #endif + cudaCheck(cudaDeviceSynchronize()); } float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, float grad_clip, int t, MultiGpuConfig* multi_gpu_config) { @@ -1105,14 +1118,14 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->m_memory, shard_num_parameters * sizeof(float))); cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float))); - cudaCheck(cudaMemset(model->m_memory, 0, shard_num_parameters * sizeof(float))); - cudaCheck(cudaMemset(model->v_memory, 0, shard_num_parameters * sizeof(float))); + cudaCheck(cudaMemsetAsync(model->m_memory, 0, shard_num_parameters * sizeof(float), model->main_stream)); + cudaCheck(cudaMemsetAsync(model->v_memory, 0, shard_num_parameters * sizeof(float), model->main_stream)); } if (model->use_master_weights == 1 && model->master_weights == NULL) { printf0("allocating %zu MiB for master copy of params\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->master_weights, shard_num_parameters * sizeof(float))); size_t grid_size = CEIL_DIV(shard_num_parameters, 512); - copy_and_cast_kernel<<>>(model->master_weights, params_memory + shard_offset, shard_num_parameters); + copy_and_cast_kernel<<main_stream>>>(model->master_weights, params_memory + shard_offset, shard_num_parameters); cudaCheck(cudaGetLastError()); } @@ -1123,11 +1136,11 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // ^1 because of the ncclReduceScatter() in gpt2_multi_gpu_grad_reduce, // grads_memory only contains the averaged gradients at the local shard // so we only calculate the grad norm at the grads_memory belonging to the local shard - global_norm_squared(grad_norm_squared, grads_memory + shard_offset, shard_num_parameters); + global_norm_squared(grad_norm_squared, grads_memory + shard_offset, shard_num_parameters, model->main_stream); } else { // the ncclAllReduce() in gpt2_multi_gpu_grad_reduce has averaged the gradients across all GPUs // so each GPU can compute the squared norm over the whole grad vector, with no added comms needed - global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters); + global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters, model->main_stream); } // transfer the gradient norm to CPU float grad_norm_squared_cpu = 0.0f; @@ -1147,9 +1160,6 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl float grad_scale = (grad_norm_cpu > grad_clip) ? grad_clip / grad_norm_cpu : 1.0f; // AdamW update - int block_size = 512; - float beta1_correction = 1.0f - powf(beta1, t); - float beta2_correction = 1.0f - powf(beta2, t); unsigned int seed = random_u32(&model->rng_state); // individually call the adamw_kernel3 on all parameter tensors separately size_t offset = 0; @@ -1200,16 +1210,16 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // - the position embeddings actively participate at every forward/backward pass float wd = (i == 0 || i == 1 || i == 4 || i == 6 || i == 10 || i == 12) ? weight_decay : 0.0f; // ok finally call the kernel - size_t num_blocks = CEIL_DIV(local_params, block_size); - adamw_kernel3<<>>(params_ptr, master_ptr, grad_ptr, - m_ptr, v_ptr, local_params, learning_rate, - beta1, beta2, beta1_correction, beta2_correction, - eps, wd, grad_scale, seed); + adamw_update(params_ptr, master_ptr, grad_ptr, + m_ptr, v_ptr, local_params, learning_rate, + beta1, beta2, t, eps, wd, grad_scale, seed, model->main_stream); } // advance the offset pointer to the next parameter tensor offset += num_parameters; } cudaCheck(cudaGetLastError()); + + cudaCheck(cudaDeviceSynchronize()); return grad_norm_cpu; } @@ -1221,10 +1231,11 @@ void gpt2_multi_gpu_param_gather(GPT2 *model, MultiGpuConfig* multi_gpu_config) // gather updated shards of model->params_memory from each process ncclCheck(ncclAllGather((floatX*)model->params_memory + multi_gpu_config->shard_offset, (floatX*)model->params_memory, multi_gpu_config->shard_num_parameters, ncclFloatX, - multi_gpu_config->nccl_comm, 0)); + multi_gpu_config->nccl_comm, model->main_stream)); } cudaCheck(cudaGetLastError()); #endif + cudaCheck(cudaDeviceSynchronize()); } float gpt2_estimate_mfu(GPT2 *model, int num_tokens, float dt) { @@ -1258,6 +1269,7 @@ void gpt2_free(GPT2 *model) { cudaCheck(cudaFree(model->grads_acts_memory)); cudaCheck(cudaFree(model->inputs)); cudaCheck(cudaFree(model->targets)); + cudaCheck(cudaStreamDestroy(model->main_stream)); cudaCheck(cudaFreeHost(model->cpu_losses)); cudaCheck(cudaFreeHost(model->cpu_losses_fp32)); free(model->workload_indices); From 4c369eea0f279e8644d482a3553036b9ddc96ff2 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 4 Jun 2024 22:47:54 +0300 Subject: [PATCH 2/4] overlap validity check and data movement --- train_gpt2.cu | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index c0731fe..526d48a 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -748,16 +748,9 @@ void gpt2_forward(GPT2 *model, const int* inputs, const int* targets, size_t B, const size_t NH = model->config.num_heads; const size_t C = model->config.channels; - // validate inputs, all indices must be in the range [0, V) - for(int i = 0; i < B * T; i++) { - assert(0 <= inputs[i] && inputs[i] < V); - if (targets != NULL) { - assert(0 <= targets[i] && targets[i] < V); - } - } - // allocate space for all the activations if needed (done here, lazily) if(model->acts_memory == NULL) { + NvtxRange rng("InitActs"); // record the current B,T as well model->batch_size = B; model->seq_len = T; @@ -792,6 +785,15 @@ void gpt2_forward(GPT2 *model, const int* inputs, const int* targets, size_t B, cudaStream_t main_stream = model->main_stream; + // validate inputs, all indices must be in the range [0, V) + // we can do this while the copies are already underway + for(int i = 0; i < B * T; i++) { + assert(0 <= inputs[i] && inputs[i] < V); + if (targets != NULL) { + assert(0 <= targets[i] && targets[i] < V); + } + } + // forward pass ParameterTensors params = model->params; // for brevity ActivationTensors acts = model->acts; @@ -912,6 +914,7 @@ void gpt2_backward(GPT2 *model, int* inputs) { // lazily allocate the memory for gradients of the weights and activations, if needed if (model->grads_memory == NULL) { + NvtxRange rng("InitGrads"); // allocate buffers for weight gradients printf0("allocating %d MiB for parameter gradients\n", (int)round(model->num_parameters * sizeof(floatX) / (1024 * 1024))); model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_elements, model->param_sizeof); @@ -1114,6 +1117,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // lazily allocate m,v memory and master weights (usually on the first iteration) if (model->m_memory == NULL) { + NvtxRange rng("InitOpt"); printf0("allocating %zu MiB for AdamW optimizer state m\n", (shard_num_parameters * sizeof(float)) >> 20); printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->m_memory, shard_num_parameters * sizeof(float))); From 0de54d10daa06829d4099cb5b862c11d890ee67a Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 5 Jun 2024 01:32:25 +0300 Subject: [PATCH 3/4] moved last remaining kernel to main stream --- llmc/encoder.cuh | 2 +- train_gpt2.cu | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/llmc/encoder.cuh b/llmc/encoder.cuh index bb3068e..1629b38 100644 --- a/llmc/encoder.cuh +++ b/llmc/encoder.cuh @@ -169,7 +169,7 @@ void encoder_backward(floatX* dwte, floatX* dwpe, floatX* scratch, // gpu output const int block_size = 256; const int N = T * C / x128::size; const int grid_size = CEIL_DIV(N, block_size); - wpe_backward_kernel<<>>(dwpe, dout, inp, B, T, C, seed); + wpe_backward_kernel<<>>(dwpe, dout, inp, B, T, C, seed); cudaCheck(cudaGetLastError()); // check the GPU scratch buffer is large enough to hold the bucket info and workload indices diff --git a/train_gpt2.cu b/train_gpt2.cu index 526d48a..37e4d23 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1674,7 +1674,7 @@ int main(int argc, char *argv[]) { int last_step = step == train_num_batches; // once in a while estimate the validation loss (all processes collaborate) - if (step % val_loss_every == 0 || last_step) { + /*if (step % val_loss_every == 0 || last_step) { NvtxRange validation_range("validation"); float val_loss = 0.0f; dataloader_reset(&val_loader); @@ -1707,7 +1707,7 @@ int main(int argc, char *argv[]) { printf0("HellaSwag: %d/%d = %f\n", (int)eval_acc_norm, eval_loader.num_examples, eval_acc_norm / eval_loader.num_examples); logger_log_eval(&logger, step, eval_acc_norm / eval_loader.num_examples); } - + */ // once in a while do model inference to print generated text (only rank 0) if (multi_gpu_config.process_rank == 0 && sample_every > 0 && (step > 0 && (step % sample_every) == 0 || last_step)) { From bc35d3438e83071208a35fde37edfe6dc8a3ad06 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 5 Jun 2024 01:50:16 +0300 Subject: [PATCH 4/4] naming the stream --- llmc/cuda_common.h | 1 + train_gpt2.cu | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/llmc/cuda_common.h b/llmc/cuda_common.h index 490dc5b..c2885e0 100644 --- a/llmc/cuda_common.h +++ b/llmc/cuda_common.h @@ -8,6 +8,7 @@ Common utilities for CUDA code. #include #include #include +#include #include #include diff --git a/train_gpt2.cu b/train_gpt2.cu index 37e4d23..83410b7 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -543,6 +543,7 @@ void gpt2_init_common(GPT2 *model) { model->use_master_weights = 1; // safe default: do keep master weights in fp32 model->recompute = 1; // good default: recompute gelu but not layernorm cudaStreamCreate(&model->main_stream); + nvtxNameCudaStreamA(model->main_stream, "main stream"); } void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { @@ -1674,7 +1675,7 @@ int main(int argc, char *argv[]) { int last_step = step == train_num_batches; // once in a while estimate the validation loss (all processes collaborate) - /*if (step % val_loss_every == 0 || last_step) { + if (step % val_loss_every == 0 || last_step) { NvtxRange validation_range("validation"); float val_loss = 0.0f; dataloader_reset(&val_loader); @@ -1707,7 +1708,7 @@ int main(int argc, char *argv[]) { printf0("HellaSwag: %d/%d = %f\n", (int)eval_acc_norm, eval_loader.num_examples, eval_acc_norm / eval_loader.num_examples); logger_log_eval(&logger, step, eval_acc_norm / eval_loader.num_examples); } - */ + // once in a while do model inference to print generated text (only rank 0) if (multi_gpu_config.process_rank == 0 && sample_every > 0 && (step > 0 && (step % sample_every) == 0 || last_step)) {