From 75d46104424e7f508b99966ed05d0ee63d69e9f0 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Mon, 24 Jun 2024 23:57:42 +0300 Subject: [PATCH] update norm computation: multi-process reduce on GPU --- llmc/global_norm.cuh | 7 ------- train_gpt2.cu | 13 +++++++------ 2 files changed, 7 insertions(+), 13 deletions(-) diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index 9e23744..e0e23b0 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -87,10 +87,3 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st global_norm_squared_kernel<<>>(out, values, count, stride); cudaCheck(cudaGetLastError()); } - -void global_norm_squared_aggregate(float* out, int max_num_block_sums, cudaStream_t stream) { - assert(max_num_block_sums > 0 && max_num_block_sums < 1024); // we need to accumulate the block sums in a single block - // important to use 1024 here for determinism, otherwise blockreduce might introduce errors - global_norm_aggregate_kernel<<<1, 1024, 0, stream>>>(out, max_num_block_sums); - cudaCheck(cudaGetLastError()); -} diff --git a/train_gpt2.cu b/train_gpt2.cu index 7127524..3f363ab 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1015,18 +1015,19 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl max_num_block_sums, is_first_pass, main_stream); } } - global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream); - cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); - // further sum the (partial) squared norm across all GPUs (see comment ^1 above) - grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu, multi_gpu_config); + global_sum_deterministic(grad_norm_squared, grad_norm_squared, max_num_block_sums, main_stream); +#if MULTI_GPU + // further sum the (partial) squared norm across all GPUs + ncclCheck(ncclAllReduce(grad_norm_squared, grad_norm_squared, sizeof(float), ncclFloat, ncclAvg, multi_gpu_config->nccl_comm, main_stream)); +#endif } else { // in regular DDP, backward 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, 0, 1, max_num_block_sums, true, main_stream); - global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream); - cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); + global_sum_deterministic(grad_norm_squared, grad_norm_squared, max_num_block_sums, main_stream); } + cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); if(!isfinite(grad_norm_squared_cpu)) { // may happen due to some issue (e.g. overflow?) // TODO: later may want to keep a global counter of instabilities like this