From 4b41087c640f1f7faec0d8a3199559d7cf2eb11a Mon Sep 17 00:00:00 2001 From: lancerts Date: Thu, 6 Jun 2024 11:46:40 -0700 Subject: [PATCH 01/20] fix the compiler warnings and errors --- dev/cuda/matmul_backward_bias.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dev/cuda/matmul_backward_bias.cu b/dev/cuda/matmul_backward_bias.cu index 16172bc..af87a69 100644 --- a/dev/cuda/matmul_backward_bias.cu +++ b/dev/cuda/matmul_backward_bias.cu @@ -2,7 +2,7 @@ Kernels for matmul backward pass bias only. Compile example: -nvcc -O3 -lcublas -lcublasLt matmul_backward_bias.cu -lineinfo -o matmul_backward_bias +nvcc -O3 -lcublas -lcublasLt -std=c++17 matmul_backward_bias.cu -lineinfo -o matmul_backward_bias ./matmul_backward_bias 1 ./matmul_backward_bias 2 @@ -116,7 +116,7 @@ __global__ void matmul_backward_bias_kernel2(floatX* dbias, const floatX* dout, sum = cg::reduce(warp, sum, cg::plus{}); // write the result to output (global memory) if(warp.thread_rank() == 0) { - dbias[idx] += sum; + dbias[idx] = (float)dbias[idx] + sum; } } @@ -148,7 +148,7 @@ __global__ void matmul_backward_bias_kernel3(floatX* dbias, const floatX* dout, float block_sum = cg::reduce(warp, warp_sum, cg::plus{}); // sum(x) // write the result to output (global memory) if(threadIdx.x == 0) { - dbias[idx] += block_sum; + dbias[idx] = (float)dbias[idx] + block_sum; } } @@ -188,7 +188,7 @@ __global__ void matmul_backward_bias_kernel4(floatX* dbias, const floatX* dout, for (int j = 0; j < vstep; j++) { dout_sum += smem[lane_id + j * warpSize]; } - dbias[tl + lane_id] += dout_sum; + dbias[tl + lane_id] = (float)dbias[tl + lane_id] + dout_sum; } } From dc8d8de06aec252617e9ca6b7ecd7ea8ac610de1 Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Thu, 13 Jun 2024 00:56:53 +0530 Subject: [PATCH 02/20] Replaced hard-coded max float with FLT_MAX --- dev/cuda/attention_backward.cu | 2 +- dev/cuda/attention_forward.cu | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index c97dbee..f6d258d 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -68,7 +68,7 @@ void attention_forward_cpu(float* out, float* preatt, float* att, float* att_bth = att + b*NH*T*T + h*T*T + t*T; // pass 1: calculate query dot key and maxval - float maxval = -10000.0f; // TODO something better + float maxval = -FLT_MAX; for (int t2 = 0; t2 < T; t2++) { // used to be t2 <= t float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index b632b4a..8c1d290 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -98,7 +98,7 @@ void attention_forward_cpu(float* out, float* preatt, float* att, float* att_bth = att + b*NH*T*T + h*T*T + t*T; // pass 1: calculate query dot key and maxval - float maxval = -10000.0f; // TODO something better + float maxval = -FLT_MAX; for (int t2 = 0; t2 <= t; t2++) { const float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key @@ -203,7 +203,7 @@ __global__ void attention_softmax_kernel1(float* att, const float* preatt, float* att_bth = att + b*NH*T*T + h*T*T + t*T; // find maxval - float maxval = -10000.0f; // TODO something better + float maxval = -FLT_MAX for (int t2 = 0; t2 <= t; t2++) { if (preatt_bth[t2] > maxval) { maxval = preatt_bth[t2]; From e2a18d9385a6e41c06d62809d58f1012563e2234 Mon Sep 17 00:00:00 2001 From: vyom1611 Date: Thu, 13 Jun 2024 00:59:15 +0530 Subject: [PATCH 03/20] Fixed missing semi-colon --- dev/cuda/attention_forward.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index 8c1d290..66b6a1b 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -203,7 +203,7 @@ __global__ void attention_softmax_kernel1(float* att, const float* preatt, float* att_bth = att + b*NH*T*T + h*T*T + t*T; // find maxval - float maxval = -FLT_MAX + float maxval = -FLT_MAX; for (int t2 = 0; t2 <= t; t2++) { if (preatt_bth[t2] > maxval) { maxval = preatt_bth[t2]; From 2d23d1f08e7de3f82abef6b743acc6cdf21b68f5 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Fri, 7 Jun 2024 16:24:44 +0200 Subject: [PATCH 04/20] Save master weights --- train_gpt2.cu | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/train_gpt2.cu b/train_gpt2.cu index 49e94c3..092dbf1 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1188,6 +1188,8 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(cpu_buffer, model->v_memory, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + cudaCheck(cudaMemcpy(cpu_buffer, model->master_weights, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); + fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); free(cpu_buffer); fclose(state_file); } @@ -1219,6 +1221,8 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename cudaCheck(cudaMemcpy(model->m_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(model->v_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); + freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + cudaCheck(cudaMemcpy(model->master_weights, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); free(cpu_buffer); fclose(state_file); } From 269b54ae45efc268a64983379ff3bf8cd69f7856 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Fri, 7 Jun 2024 16:34:59 +0200 Subject: [PATCH 05/20] Allocate memory in load state --- train_gpt2.cu | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 092dbf1..888a0ad 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1173,6 +1173,7 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) state_header[1] = 1; // version number state_header[2] = multi_gpu_config.num_processes; // number of processes state_header[3] = multi_gpu_config.process_rank; // rank of this process + state_header[4] = model->use_master_weights; // whether we're using fp32 master weights // int main state, start at 10 to leave some padding state_header[10] = step; // step of the optimization // model state, state, start at 20 to leave some padding @@ -1188,8 +1189,10 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(cpu_buffer, model->v_memory, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); - cudaCheck(cudaMemcpy(cpu_buffer, model->master_weights, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); - fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + if (model->use_master_weights == 1) { + cudaCheck(cudaMemcpy(cpu_buffer, model->master_weights, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); + fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + } free(cpu_buffer); fclose(state_file); } @@ -1216,13 +1219,19 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename cudaCheck(cudaMalloc((void**)&model->m_memory, shard_num_parameters * sizeof(float))); cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float))); } + if (state_header[4] == 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))); + } float* cpu_buffer = (float*)mallocCheck(shard_num_parameters * sizeof(float)); freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(model->m_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(model->v_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); - freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); - cudaCheck(cudaMemcpy(model->master_weights, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); + if (state_header[4] == 1) { // if we used master weights during the state save + freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + cudaCheck(cudaMemcpy(model->master_weights, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); + } free(cpu_buffer); fclose(state_file); } From 0f144e3052e4365f376619a5c2fdba2f4b5a3a5b Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Thu, 13 Jun 2024 10:36:16 +0000 Subject: [PATCH 06/20] Support 100% determinism, fix race condition, update global norm --- llmc/fused_classifier.cuh | 6 ++++++ llmc/global_norm.cuh | 36 +++++++++++++++++++++++++++--------- train_gpt2.cu | 5 +++-- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/llmc/fused_classifier.cuh b/llmc/fused_classifier.cuh index b13fa35..a52765d 100644 --- a/llmc/fused_classifier.cuh +++ b/llmc/fused_classifier.cuh @@ -85,6 +85,12 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS) losses[idx] = (floatX)(-logf(prob)); } + // without this synchronization point we have a race condition: + // the logits used above to compute the loss are concurrently (race) modified to carry backward pass grads. + // since the "logits" are overwritten to be in the [-1, 1] range and sp.Offset is sometimes smaller than -90 + // we errouneously end up computing exp^(90+) which gives us infinities in the loss! this is the fix. + __syncthreads(); + // calculate the gradients directly, saves bandwidth from probs during training // but also supports writing probs for inference-only and debugging const floatX* logits_vec = logits + idx * P; diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index 92866e8..751224d 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -13,13 +13,7 @@ Global norm, used in gradient clipping template __device__ float global_norm_squared_for_range(const T* data, size_t count) { - // we want as few atomics as possible, so each block tries to do - // the maximum amount of work (so no fixed chunk, but instead iterating - // until we run out of data), and then we reduce inside the block - // and finally have just one atomic per block. - // out will be updated atomically from all thread blocks. It is a float, so the - // atomic op is unproblematic - size_t index = threadIdx.x + blockDim.x * blockIdx.x; + size_t index = blockIdx.x * blockDim.x + threadIdx.x; size_t grid_width = blockDim.x * gridDim.x; float accumulator = 0.f; for(size_t i = index; i < count; i += grid_width) { @@ -32,8 +26,22 @@ __device__ float global_norm_squared_for_range(const T* data, size_t count) { template __global__ void global_norm_squared_kernel(float* out, const T* data, size_t count, ptrdiff_t stride) { float block_sum = global_norm_squared_for_range(data + blockIdx.y * stride, count); + // each block accumulates its partial sum to out[out_index] + // we want to avoid using atomic add here so we combine this kernel with another kernel call + // that sums up the partial block sums if(threadIdx.x == 0) { - atomicAdd(out, block_sum); + size_t out_index = blockIdx.y * gridDim.x + blockIdx.x; + *(out + out_index) = *(out + out_index) + block_sum; + } +} + +__global__ void global_norm_aggregate_kernel(float* out, size_t grid_size) { + size_t index = threadIdx.x; + // grab block sums from the previous kernel, use 0. as the neutral sum element + float block_sum = (index < grid_size) ? out[index] : 0.f; + float sum = blockReduce(block_sum); + if(threadIdx.x == 0) { + *out = sum; // out[0] ends up with the final norm squared } } @@ -50,9 +58,10 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st // on all gpus, so the division really is going to be exact. const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; assert(grid_size > 0); // gives a better error than letting the call below fail + assert(grid_size < 1024); // we want to later accumulate the block sums in a single block // initialize out with zero if(reset) { - cudaCheck(cudaMemsetAsync(out, 0, sizeof(float), stream)); + cudaCheck(cudaMemsetAsync(out, 0, grid_size * sizeof(float), stream)); } const int gx = CEIL_DIV(grid_size, num_slices); const int gy = num_slices; @@ -60,3 +69,12 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st cudaCheck(cudaGetLastError()); } +void global_norm_squared_aggregate(float* out, cudaStream_t stream) { + const int block_size = 512; + const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; + assert(grid_size > 0); + assert(grid_size < 1024); + + global_norm_aggregate_kernel<<<1, grid_size, 0, stream>>>(out, grid_size); + cudaCheck(cudaGetLastError()); +} \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index 888a0ad..37a97c9 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -974,13 +974,12 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // because of the ncclReduceScatter() in backward, // grads_memory only contains the averaged gradients at the local shards, // so we only calculate the grad norm at the grads_memory belonging to the local shards - cudaCheck(cudaMemsetAsync(grad_norm_squared, 0, sizeof(float), main_stream)); for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { if((i < 2 || i > 13)) { ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i); ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); ptrdiff_t offset = tensor.offset + shard.offset; - global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, false, main_stream); + global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, i == 0, main_stream); } else { ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i); ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); @@ -989,6 +988,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl false, main_stream); } } + global_norm_squared_aggregate(grad_norm_squared, 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); @@ -996,6 +996,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // 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, true, main_stream); + global_norm_squared_aggregate(grad_norm_squared, main_stream); cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); } From c8cd02306344d0bee9d39f38bf630cadfddb5db2 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Fri, 14 Jun 2024 21:44:20 +0000 Subject: [PATCH 07/20] Full determinism --- llmc/global_norm.cuh | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index 751224d..d87ba30 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -31,7 +31,7 @@ __global__ void global_norm_squared_kernel(float* out, const T* data, size_t cou // that sums up the partial block sums if(threadIdx.x == 0) { size_t out_index = blockIdx.y * gridDim.x + blockIdx.x; - *(out + out_index) = *(out + out_index) + block_sum; + out[out_index] = out[out_index] + block_sum; } } @@ -41,7 +41,7 @@ __global__ void global_norm_aggregate_kernel(float* out, size_t grid_size) { float block_sum = (index < grid_size) ? out[index] : 0.f; float sum = blockReduce(block_sum); if(threadIdx.x == 0) { - *out = sum; // out[0] ends up with the final norm squared + out[0] = sum; // out[0] ends up with the final norm squared } } @@ -59,12 +59,15 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; assert(grid_size > 0); // gives a better error than letting the call below fail assert(grid_size < 1024); // we want to later accumulate the block sums in a single block - // initialize out with zero - if(reset) { - cudaCheck(cudaMemsetAsync(out, 0, grid_size * sizeof(float), stream)); - } + // if not the case we have to find the biggest gx*gy and pass that into `global_norm_squared_aggregate` + assert(grid_size % num_slices == 0); + const int gx = CEIL_DIV(grid_size, num_slices); const int gy = num_slices; + + if (reset) { + cudaCheck(cudaMemsetAsync(out, 0, gx * gy * sizeof(float), stream)); + } global_norm_squared_kernel<<>>(out, values, count, stride); cudaCheck(cudaGetLastError()); } @@ -72,9 +75,7 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st void global_norm_squared_aggregate(float* out, cudaStream_t stream) { const int block_size = 512; const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; - assert(grid_size > 0); - assert(grid_size < 1024); - - global_norm_aggregate_kernel<<<1, grid_size, 0, stream>>>(out, grid_size); + assert(grid_size > 0 && grid_size < 1024); // we need to accumulate the block sums in a single block + global_norm_aggregate_kernel<<<1, 1024, 0, stream>>>(out, grid_size); cudaCheck(cudaGetLastError()); } \ No newline at end of file From 20f0e7f4e89da0a7b230612fb0694bf2b98bb328 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 15 Jun 2024 08:15:38 +0000 Subject: [PATCH 08/20] Fix copy and casting from spawning excessive threads --- train_gpt2.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 49e94c3..c30757e 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1037,7 +1037,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl float* master_ptr = NULL; if (model->master_weights != NULL) { master_ptr = model->master_weights + opt_state_offset; } if(init_master_weights) { - size_t grid_size = CEIL_DIV(shard_num_parameters, 512); + size_t grid_size = CEIL_DIV(shard.size, 512); copy_and_cast_kernel<<>>(master_ptr, param_ptr, shard.size, shard.size, tensor.size); cudaCheck(cudaGetLastError()); From ff14b0819ff4a7910f597c246ee7351e538442fb Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 15 Jun 2024 12:02:39 +0000 Subject: [PATCH 09/20] Make sure each parameter was rounded with a unique seed --- llmc/adamw.cuh | 4 +--- llmc/cuda_utils.cuh | 3 ++- train_gpt2.cu | 4 +++- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index 94d4642..f806eaa 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -40,9 +40,7 @@ __device__ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* float param = old_param - (learning_rate * (m / (sqrtf(v) + eps) + weight_decay * old_param)); // update our low precision version of the parameters using stochastic rounding // this will be used in the next forward pass - // TODO: simply doing `params_memory[i] = (floatX)param;` breaks everything (why?) - unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x + blockDim.y * gridDim.y, seed); - stochastic_rounding(param, ¶ms_memory[idx], random); + stochastic_rounding(param, ¶ms_memory[idx], seed); // write the full, float version of the param into our master copy, if we maintain one // this will be used in the next update if (master_params_memory != NULL) { master_params_memory[idx] = param; } diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 1bca60a..fe3b265 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -190,7 +190,8 @@ __device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY // stochastic rounding built on top of Squirel Noise above (with seed updated per step via xorshift) __device__ __forceinline__ void stochastic_rounding(float in, __nv_bfloat16 *out, unsigned int seed) { // todo - is this stochastic rounding *too good*? can we cut any corners? - unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); + // makes sure each thread gets a different random number + unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x * blockDim.x + blockIdx.y, seed); unsigned int threshold = random & 0xFFFF; unsigned int float_bits = __float_as_uint(in); unsigned int rounded_bits = float_bits & 0x0000FFFF; diff --git a/train_gpt2.cu b/train_gpt2.cu index 49e94c3..40437d2 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1009,10 +1009,12 @@ 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 - unsigned int seed = random_u32(&model->rng_state); // handle adamw for all the transformer blocks for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + // generate a unique seed for each tensor + unsigned int seed = random_u32(&model->rng_state); + int num_layers = model->config.num_layers; if((i < 2 || i > 13)) { num_layers = 1; From 0f6853d39970b11bd2e6d2d29eb088c4c18f677d Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 15 Jun 2024 17:29:14 +0200 Subject: [PATCH 10/20] Add global norm prod kernels to dev --- dev/cuda/common.h | 27 ++++++++++++ dev/cuda/global_norm.cu | 94 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 117 insertions(+), 4 deletions(-) diff --git a/dev/cuda/common.h b/dev/cuda/common.h index 6502baa..b08f36b 100644 --- a/dev/cuda/common.h +++ b/dev/cuda/common.h @@ -5,6 +5,7 @@ #include #include +#define WARP_SIZE 32U template __host__ __device__ T ceil_div(T dividend, T divisor) { @@ -18,6 +19,32 @@ __device__ float warpReduceSum(float val) { return val; } +// requires all 32 threads in the warp to be active, but should work for any block size +// uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes +// the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end +// but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1 +using reduction_func_t = float (*) (float); +template +__device__ inline float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) { + // two reductions of up to 1024 threads: + // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) + __shared__ float shared_val[WARP_SIZE]; + const int lane_id = threadIdx.x % WARP_SIZE; + const int warp_id = threadIdx.x / WARP_SIZE; + const int num_warps = blockDim.x / WARP_SIZE; + + float warp_val = warp_reduction(val); + if (lane_id == 0) { shared_val[warp_id] = warp_val; } + __syncthreads(); + warp_val = (lane_id < num_warps) ? shared_val[lane_id] : out_of_bounds; + float block_val = warp_reduction(warp_val); + + if (final_sync) { + __syncthreads(); // only needed in loops when effectively reusing shared memory etc. + } + return block_val; +} + // ---------------------------------------------------------------------------- // checking utils diff --git a/dev/cuda/global_norm.cu b/dev/cuda/global_norm.cu index 6c2ed03..bbdaca1 100644 --- a/dev/cuda/global_norm.cu +++ b/dev/cuda/global_norm.cu @@ -89,6 +89,54 @@ __global__ void norm_kernel2(float* out, const T* data, size_t count) { } } +template +__global__ void norm_kernel3(float* out, const T* data, size_t count) { + size_t index = blockIdx.x * blockDim.x + threadIdx.x; + size_t grid_width = blockDim.x * gridDim.x; + float accumulator = 0.f; + for(size_t i = index; i < count; i += grid_width) { + accumulator += (float)data[i] * (float)data[i]; + } + // block-level reduce + float block_sum = blockReduce(accumulator); + if(threadIdx.x == 0) { + atomicAdd(out, block_sum); + } +} + +// Same as kernel3 but without atomic adds -> this allows us to have determinism due to the +// non associativity of floating point operations. Roughly same performance as kernel3. +template +__global__ void norm_kernel4(float* out, const T* data, size_t count) { + size_t index = blockIdx.x * blockDim.x + threadIdx.x; + size_t grid_width = blockDim.x * gridDim.x; + float accumulator = 0.f; + for(size_t i = index; i < count; i += grid_width) { + accumulator += (float)data[i] * (float)data[i]; + } + // block-level reduce + float block_sum = blockReduce(accumulator); + // each block accumulates its partial sum to out[blockIdx.x] + // we want to avoid using atomic add here so we combine this kernel with the aggregate kernel call + // that sums up the partial block sums + if(threadIdx.x == 0) { + out[blockIdx.x] = block_sum; + } +} + +__global__ void global_norm_aggregate_kernel(float* out, size_t count) { + size_t index = threadIdx.x; + // grab block sums from the previous kernel, use 0. as the neutral sum element + float block_sum = (index < count) ? out[index] : 0.f; + float sum = blockReduce(block_sum); + if(threadIdx.x == 0) { + out[0] = sum; // out[0] ends up with the final norm squared + } +} + +// ---------------------------------------------------------------------------- +// kernel launchers + template void global_norm1(float* out, const T* values, size_t count, int block_size) { // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. @@ -111,17 +159,55 @@ void global_norm2(float* out, const T* values, size_t count, int block_size) { cudaCheck(cudaGetLastError()); } -void global_norm(int kernel_num, float* out, const floatX* values, size_t count, int block_size) { +template +void global_norm3(float* out, const T* values, size_t count, int block_size, cudaDeviceProp deviceProp) { + // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. + // having one block less than possible is a tiny performance hit, having + // one block too many is catastrophic, since it only can start once all the other + // blocks finish. anyway, I think cuda_threads_per_SM should be a multiple of 512 + // on all gpus, so the division really is going to be exact. + const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + norm_kernel3<<>>(out, values, count); + cudaCheck(cudaGetLastError()); +} + +template +void global_norm4(float* out, const T* values, size_t count, int block_size, cudaDeviceProp deviceProp) { + if (block_size <= 64) { + block_size = 128; // to avoid triggering the assert below + } + // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. + // having one block less than possible is a tiny performance hit, having + // one block too many is catastrophic, since it only can start once all the other + // blocks finish. anyway, I think cuda_threads_per_SM should be a multiple of 512 + // on all gpus, so the division really is going to be exact. + const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + assert(grid_size < 1024); // we want to later accumulate the block sums in a single block + norm_kernel4<<>>(out, values, count); + cudaCheck(cudaGetLastError()); + global_norm_aggregate_kernel<<<1, 1024>>>(out, grid_size); + cudaCheck(cudaGetLastError()); +} + +void global_norm(int kernel_num, float* out, const floatX* values, size_t count, int block_size, cudaDeviceProp deviceProp) { switch (kernel_num) { case 1: return global_norm1(out, values, count, block_size); case 2: return global_norm2(out, values, count, block_size); + case 3: + return global_norm3(out, values, count, block_size, deviceProp); + case 4: + return global_norm4(out, values, count, block_size, deviceProp); } } int main(int argc, const char **argv) { setup_main(); + cudaDeviceProp deviceProp; + cudaGetDeviceProperties(&deviceProp, 0); int C = 768; int L = 12; @@ -148,7 +234,7 @@ int main(int argc, const char **argv) { // move to GPU float* d_out; floatX* d_inp; - cudaCheck(cudaMalloc(&d_out, sizeof(float))); + cudaCheck(cudaMalloc(&d_out, 1024 * sizeof(float))); // 1024 needed for kernel 4 cudaCheck(cudaMalloc(&d_inp, num_params * sizeof(floatX))); cudaCheck(memcpy_convert(d_inp, inp, num_params)); @@ -157,7 +243,7 @@ int main(int argc, const char **argv) { int block_size = block_sizes[j]; printf("Checking block size %d.\n", block_size); cudaCheck(cudaMemset(d_out, 0, sizeof(float))); - global_norm(kernel_num, d_out, d_inp, num_params, block_size); + global_norm(kernel_num, d_out, d_inp, num_params, block_size, deviceProp); validate_result(d_out, &out, "out", 1, 1e-2f); } @@ -170,7 +256,7 @@ int main(int argc, const char **argv) { float elapsed_time = benchmark_kernel(repeat_times, global_norm, kernel_num, d_out, d_inp, - num_params, block_size); + num_params, block_size, deviceProp); size_t memory_ops = num_params * sizeof(floatX); float memory_bandwidth = memory_ops / elapsed_time / 1e6; From f3195d19afe37e798bc25454bbd855c525b8dec9 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 15 Jun 2024 17:40:14 +0200 Subject: [PATCH 11/20] Expose deviceProp as a global var --- dev/cuda/common.h | 1 + dev/cuda/global_norm.cu | 16 ++++++++-------- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/dev/cuda/common.h b/dev/cuda/common.h index b08f36b..9ad8578 100644 --- a/dev/cuda/common.h +++ b/dev/cuda/common.h @@ -6,6 +6,7 @@ #include #define WARP_SIZE 32U +extern cudaDeviceProp deviceProp; template __host__ __device__ T ceil_div(T dividend, T divisor) { diff --git a/dev/cuda/global_norm.cu b/dev/cuda/global_norm.cu index bbdaca1..f54a35a 100644 --- a/dev/cuda/global_norm.cu +++ b/dev/cuda/global_norm.cu @@ -16,6 +16,7 @@ nvcc -O3 --use_fast_math global_norm.cu -o global_norm #define ENABLE_BF16 #include "common.h" +cudaDeviceProp deviceProp; float global_norm_cpu(const float* data, size_t count) { // accumulate in double so we have an accurate numerical reference @@ -160,7 +161,7 @@ void global_norm2(float* out, const T* values, size_t count, int block_size) { } template -void global_norm3(float* out, const T* values, size_t count, int block_size, cudaDeviceProp deviceProp) { +void global_norm3(float* out, const T* values, size_t count, int block_size) { // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. // having one block less than possible is a tiny performance hit, having // one block too many is catastrophic, since it only can start once all the other @@ -173,7 +174,7 @@ void global_norm3(float* out, const T* values, size_t count, int block_size, cud } template -void global_norm4(float* out, const T* values, size_t count, int block_size, cudaDeviceProp deviceProp) { +void global_norm4(float* out, const T* values, size_t count, int block_size) { if (block_size <= 64) { block_size = 128; // to avoid triggering the assert below } @@ -191,22 +192,21 @@ void global_norm4(float* out, const T* values, size_t count, int block_size, cud cudaCheck(cudaGetLastError()); } -void global_norm(int kernel_num, float* out, const floatX* values, size_t count, int block_size, cudaDeviceProp deviceProp) { +void global_norm(int kernel_num, float* out, const floatX* values, size_t count, int block_size) { switch (kernel_num) { case 1: return global_norm1(out, values, count, block_size); case 2: return global_norm2(out, values, count, block_size); case 3: - return global_norm3(out, values, count, block_size, deviceProp); + return global_norm3(out, values, count, block_size); case 4: - return global_norm4(out, values, count, block_size, deviceProp); + return global_norm4(out, values, count, block_size); } } int main(int argc, const char **argv) { setup_main(); - cudaDeviceProp deviceProp; cudaGetDeviceProperties(&deviceProp, 0); int C = 768; @@ -243,7 +243,7 @@ int main(int argc, const char **argv) { int block_size = block_sizes[j]; printf("Checking block size %d.\n", block_size); cudaCheck(cudaMemset(d_out, 0, sizeof(float))); - global_norm(kernel_num, d_out, d_inp, num_params, block_size, deviceProp); + global_norm(kernel_num, d_out, d_inp, num_params, block_size); validate_result(d_out, &out, "out", 1, 1e-2f); } @@ -256,7 +256,7 @@ int main(int argc, const char **argv) { float elapsed_time = benchmark_kernel(repeat_times, global_norm, kernel_num, d_out, d_inp, - num_params, block_size, deviceProp); + num_params, block_size); size_t memory_ops = num_params * sizeof(floatX); float memory_bandwidth = memory_ops / elapsed_time / 1e6; From 0f112f1b6225b28e3cf827dc658b115a8f6a8012 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 15 Jun 2024 19:39:05 +0200 Subject: [PATCH 12/20] Fix CI complaining about default args --- dev/cuda/common.h | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/dev/cuda/common.h b/dev/cuda/common.h index 9ad8578..61a783a 100644 --- a/dev/cuda/common.h +++ b/dev/cuda/common.h @@ -25,8 +25,9 @@ __device__ float warpReduceSum(float val) { // the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end // but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1 using reduction_func_t = float (*) (float); + template -__device__ inline float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) { +__device__ inline float blockReduce(float val, bool final_sync, float out_of_bounds) { // two reductions of up to 1024 threads: // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) __shared__ float shared_val[WARP_SIZE]; @@ -46,6 +47,12 @@ __device__ inline float blockReduce(float val, bool final_sync=false, float out_ return block_val; } +// Helper function to call blockReduce with default arguments +template +__device__ inline float blockReduce(float val) { + return blockReduce(val, false, 0.0f); +} + // ---------------------------------------------------------------------------- // checking utils From 09910acc0bd7b3c4270656a97bffa1ee3f5f8ef1 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 15 Jun 2024 19:55:11 +0200 Subject: [PATCH 13/20] Remove blockreduce from classifier fused, move to common.h --- dev/cuda/classifier_fused.cu | 27 --------------------------- 1 file changed, 27 deletions(-) diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index 5b6c986..d58f101 100644 --- a/dev/cuda/classifier_fused.cu +++ b/dev/cuda/classifier_fused.cu @@ -481,33 +481,6 @@ __global__ void fused_classifier_kernel4(floatX* dlogits, floatX* losses, floatX } } -// todo - move to common.h - or ideally somewhere it's not duplicated between train & common? -// requires all 32 threads in the warp to be active, but should work for any block size -// uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes -// the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end -// but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1 -using reduction_func_t = float (*) (float); -template -__device__ float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) { - // two reductions of up to 1024 threads: - // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) - __shared__ float shared_val[32]; - const int lane_id = threadIdx.x % 32; - const int warp_id = threadIdx.x / 32; - const int num_warps = blockDim.x / 32; - - float warp_val = warp_reduction(val); - if (lane_id == 0) { shared_val[warp_id] = warp_val; } - __syncthreads(); - warp_val = (lane_id < num_warps) ? shared_val[lane_id] : out_of_bounds; - float block_val = warp_reduction(warp_val); - - if (final_sync) { - __syncthreads(); // only needed in loops when effectively reusing shared memory etc. - } - return block_val; -} - __device__ SoftmaxParams prepare_softmax_blockwide3(int64_t idx, const floatX* inp, int V, int P) { // same but not float4 // one row of inp, i.e. inp[idx, :] of shape (V,) From d003274b7cc76b1da1e2fb300d51e23d33db4dc5 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 15 Jun 2024 20:23:59 +0200 Subject: [PATCH 14/20] Remove warp size - moved to common.h --- dev/cuda/layernorm_backward.cu | 2 -- 1 file changed, 2 deletions(-) diff --git a/dev/cuda/layernorm_backward.cu b/dev/cuda/layernorm_backward.cu index dc9d7e9..3930cec 100644 --- a/dev/cuda/layernorm_backward.cu +++ b/dev/cuda/layernorm_backward.cu @@ -874,7 +874,6 @@ __global__ void layernorm_backward_kernel9(floatX* dinp, floatX* dweight, floatX } __trap(); // prefer to crash here than run into a deadlock later on } - constexpr int WARP_SIZE = 32; int BLOCK_SIZE = blockDim.x; int warpsInBlock = BLOCK_SIZE / WARP_SIZE; //number of warps in block extern __shared__ float shared[]; // size = 2 * C + 1 @@ -1059,7 +1058,6 @@ layernorm_backward_kernel10(floatX* dinp, floatX* dweight, floatX* dbias, float* const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { - constexpr int WARP_SIZE = 32; int BLOCK_SIZE = blockDim.x; int warpsInBlock = BLOCK_SIZE / WARP_SIZE; //number of warps in block extern __shared__ float shared[]; // size = 2 * C + 1 From d5601d16da95c7bb825a982154b829ec46eeb326 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 16 Jun 2024 09:34:01 +0200 Subject: [PATCH 15/20] Fix stuff in layernorm forward --- dev/cuda/layernorm_forward.cu | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/dev/cuda/layernorm_forward.cu b/dev/cuda/layernorm_forward.cu index 3e94828..8631aa1 100644 --- a/dev/cuda/layernorm_forward.cu +++ b/dev/cuda/layernorm_forward.cu @@ -290,7 +290,7 @@ __global__ void layernorm_forward_kernel5(float* __restrict__ out, float* __rest int num_warps = blockDim.x / 32; int warp_id = threadIdx.x / 32; int lane_id = threadIdx.x % 32; - int idx = blockIdx.x; // simpoy one block per row + int idx = blockIdx.x; // simply one block per row // the row of input that this group of threads is responsible for const float* x = inp + idx * C; // thread coarsening through the row, reduce the sum in series @@ -356,9 +356,9 @@ void layernorm_forward2(float* out, float* mean, float* rstd, const int block_size) { int N = B * T; // in mean and rstd, threads cooperate within blocks via reductions - mean_kernel<<>>(mean, inp, N, C, block_size); + mean_kernel<<>>(mean, inp, N, C, block_size); cudaCheck(cudaGetLastError()); - rstd_kernel<<>>(rstd, inp, mean, N, C, block_size); + rstd_kernel<<>>(rstd, inp, mean, N, C, block_size); cudaCheck(cudaGetLastError()); // in the normalization, everything just gets flattened out const int block_size2 = 256; @@ -394,6 +394,7 @@ void layernorm_forward5(float* out, float* mean, float* rstd, int B, int T, int C, const int block_size) { assert(block_size % 32 == 0); + assert(block_size <= 1024); const int N = B * T; const int grid_size = N; layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); @@ -473,9 +474,6 @@ int main(int argc, char **argv) { printf("Using kernel %d\n", kernel_num); int block_sizes[] = {32, 64, 128, 256, 512, 1024}; - float* out_gpu = (float*)malloc(B * T * C * sizeof(float)); - float* mean_gpu = (float*)malloc(B * T * sizeof(float)); - float* rstd_gpu = (float*)malloc(B * T * sizeof(float)); layernorm_forward_cpu(out, mean, rstd, inp, weight, bias, B, T, C); From 9caf2c7c2f99c3c9d94537e21296a69d67394520 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 16 Jun 2024 23:49:34 +0200 Subject: [PATCH 16/20] Proper handling of max num of block sums --- llmc/global_norm.cuh | 32 +++++++++++++++++++++++--------- train_gpt2.cu | 12 +++++++----- 2 files changed, 30 insertions(+), 14 deletions(-) diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index d87ba30..26bcf7e 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -49,7 +49,7 @@ __global__ void global_norm_aggregate_kernel(float* out, size_t grid_size) { // kernel launcher template -void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t stride, int num_slices, bool reset, cudaStream_t stream) { +void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t stride, int num_slices, int max_num_block_sums, bool reset, 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 @@ -58,24 +58,38 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st // on all gpus, so the division really is going to be exact. const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; assert(grid_size > 0); // gives a better error than letting the call below fail - assert(grid_size < 1024); // we want to later accumulate the block sums in a single block - // if not the case we have to find the biggest gx*gy and pass that into `global_norm_squared_aggregate` - assert(grid_size % num_slices == 0); const int gx = CEIL_DIV(grid_size, num_slices); const int gy = num_slices; + assert(gx * gy < 1024); // we want to later accumulate the block sums in a single block + if (reset) { - cudaCheck(cudaMemsetAsync(out, 0, gx * gy * sizeof(float), stream)); + cudaCheck(cudaMemsetAsync(out, 0, max_num_block_sums * sizeof(float), stream)); } global_norm_squared_kernel<<>>(out, values, count, stride); cudaCheck(cudaGetLastError()); } -void global_norm_squared_aggregate(float* out, cudaStream_t stream) { +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()); +} + +// Helper function determines the maximum number of block sums +int get_max_num_block_sums(int* num_slices_all, int numel) { const int block_size = 512; const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; - assert(grid_size > 0 && grid_size < 1024); // we need to accumulate the block sums in a single block - global_norm_aggregate_kernel<<<1, 1024, 0, stream>>>(out, grid_size); - cudaCheck(cudaGetLastError()); + assert(grid_size > 0); + int max_num_block_sums = 0; + for (int i = 0; i < numel; i++) { + int num_slices = num_slices_all[i]; + const int gx = CEIL_DIV(grid_size, num_slices); + const int gy = num_slices; + max_num_block_sums = max(max_num_block_sums, gx * gy); + } + + return max_num_block_sums; } \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index 9556189..39bfa2f 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -970,6 +970,8 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl float* grad_norm_squared = (float*)model->acts.output; float grad_norm_squared_cpu = 0.0f; + int num_slices[2] = {1, model->config.num_layers}; + int max_num_block_sums = get_max_num_block_sums(num_slices, 2); if (multi_gpu_config->zero_stage == 1) { // because of the ncclReduceScatter() in backward, // grads_memory only contains the averaged gradients at the local shards, @@ -979,24 +981,24 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i); ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); ptrdiff_t offset = tensor.offset + shard.offset; - global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, i == 0, main_stream); + global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, max_num_block_sums, i == 0, main_stream); } else { ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i); ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); ptrdiff_t offset = tensor.offset + shard.offset; global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, tensor.size, model->config.num_layers, - false, main_stream); + max_num_block_sums, false, main_stream); } } - global_norm_squared_aggregate(grad_norm_squared, 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); } 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, true, main_stream); - global_norm_squared_aggregate(grad_norm_squared, main_stream); + 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)); } From 8f724dbb72cd3397298d0013773d817f9e371488 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 16 Jun 2024 23:54:09 +0200 Subject: [PATCH 17/20] Add a minor note that code is coupled --- llmc/global_norm.cuh | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index 26bcf7e..9e23744 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -48,6 +48,23 @@ __global__ void global_norm_aggregate_kernel(float* out, size_t grid_size) { // ---------------------------------------------------------------------------- // kernel launcher +// Helper function determines the maximum number of block sums +int get_max_num_block_sums(int* num_slices_all, int numel) { + // NOTE: this needs to be kept in sync with `global_norm_squared` below. + const int block_size = 512; + const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; + assert(grid_size > 0); + int max_num_block_sums = 0; + for (int i = 0; i < numel; i++) { + int num_slices = num_slices_all[i]; + const int gx = CEIL_DIV(grid_size, num_slices); + const int gy = num_slices; + max_num_block_sums = max(max_num_block_sums, gx * gy); + } + + return max_num_block_sums; +} + template void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t stride, int num_slices, int max_num_block_sums, bool reset, cudaStream_t stream) { const int block_size = 512; @@ -77,19 +94,3 @@ void global_norm_squared_aggregate(float* out, int max_num_block_sums, cudaStrea global_norm_aggregate_kernel<<<1, 1024, 0, stream>>>(out, max_num_block_sums); cudaCheck(cudaGetLastError()); } - -// Helper function determines the maximum number of block sums -int get_max_num_block_sums(int* num_slices_all, int numel) { - const int block_size = 512; - const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; - assert(grid_size > 0); - int max_num_block_sums = 0; - for (int i = 0; i < numel; i++) { - int num_slices = num_slices_all[i]; - const int gx = CEIL_DIV(grid_size, num_slices); - const int gy = num_slices; - max_num_block_sums = max(max_num_block_sums, gx * gy); - } - - return max_num_block_sums; -} \ No newline at end of file From 0eec8cee4d714f38a802d95736b297f11ffbeedf Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 16 Jun 2024 23:59:29 +0200 Subject: [PATCH 18/20] Patch numpy version --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 80471a8..ea4bc76 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ tqdm -numpy +numpy<2 torch tiktoken transformers From a0a8795be111b4123f8a57c7e20aba017e83a756 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 14:12:05 +0200 Subject: [PATCH 19/20] Fix guarding logic, refactor --- .github/workflows/ci_gpu.yml | 22 +++++++++++----------- dev/cuda/adamw.cu | 4 ++-- dev/cuda/classifier_fused.cu | 7 ++++--- dev/cuda/crossentropy_softmax_backward.cu | 2 +- llmc/gelu.cuh | 4 ++-- llmc/layernorm.cuh | 2 +- 6 files changed, 21 insertions(+), 20 deletions(-) diff --git a/.github/workflows/ci_gpu.yml b/.github/workflows/ci_gpu.yml index d151506..9162d3b 100644 --- a/.github/workflows/ci_gpu.yml +++ b/.github/workflows/ci_gpu.yml @@ -31,19 +31,19 @@ jobs: run: python train_gpt2.py - name: Compile training and testing program - run: make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu + run: make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu - name: Train model (With OpenMP) run: OMP_NUM_THREADS=8 ./train_gpt2cu - name: Train model (FP32) with gpt2_124M.bin - run: | + run: | PRECISION=FP32 make train_gpt2cu ./train_gpt2cu -b 4 -t 64 -l 1e-4 -v 200 -s 200 -a 1 -x 10 -e gpt2_124M.bin - + - name: Build FP32 precision run: PRECISION=FP32 make test_gpt2cu profile_gpt2cu - + - name: Run default run: ./test_gpt2cu @@ -52,7 +52,7 @@ jobs: - name: Run recompute LN run: ./test_gpt2cu -r 2 - + - name: Build BF16 precision run: PRECISION=BF16 make train_gpt2cu test_gpt2cu profile_gpt2cu @@ -67,18 +67,18 @@ jobs: - name: Run recompute LN run: ./test_gpt2cu -r 2 - + - name: Train model fp32 (With OpenMP) run: OMP_NUM_THREADS=8 ./train_gpt2fp32cu - name: Execute testing program (With OpenMP) run: OMP_NUM_THREADS=8 ./test_gpt2cu - + - name: Execute testing program fp32 (With OpenMP) run: OMP_NUM_THREADS=8 ./test_gpt2fp32cu - name: Compile training and testing program without OpenMP - run: NO_OMP=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu + run: NO_OMP=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu - name: Train model (No OpenMP) run: NO_OMP=1 ./train_gpt2cu @@ -88,14 +88,14 @@ jobs: - name: Execute testing program (No OpenMP) run: ./test_gpt2cu -b 32 - + - name: Execute testing program fp32 (No OpenMP) run: ./test_gpt2fp32cu - name: Install cuDNN-frontend - run: + run: git clone https://github.com/NVIDIA/cudnn-frontend.git - + - name: Build with cuDNN run: USE_CUDNN=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu diff --git a/dev/cuda/adamw.cu b/dev/cuda/adamw.cu index 20a6560..f27600a 100644 --- a/dev/cuda/adamw.cu +++ b/dev/cuda/adamw.cu @@ -159,8 +159,8 @@ int main(int argc, char **argv) { // create random data on host (to be used for the CPU reference implementation) float* params_memory = make_random_float(num_parameters); float* grads_memory = make_random_float(num_parameters); - float* m_memory = make_random_float_01(num_parameters); - float* v_memory = make_random_float_01(num_parameters); + float* m_memory = make_random_float(num_parameters); + float* v_memory = make_random_float(num_parameters); // move to GPU float* d_params_memory; diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index d58f101..9d93f9d 100644 --- a/dev/cuda/classifier_fused.cu +++ b/dev/cuda/classifier_fused.cu @@ -114,7 +114,7 @@ __device__ SoftmaxParams prepare_softmax(cg::thread_block_tile<32>& warp, int64_t idx, const float* inp, int V, int P) { // this warp (of 32) threads processes one row of inp, i.e. inp[idx, :] of shape (V,) // note that inp is actually (B * T, P) but we only use the first V elements - // this function tehen calculates: + // this function then calculates: // 1) the max value to subtract for numerical stability and // 2) the sum normalization factor const float* x = inp + idx * P; @@ -680,8 +680,8 @@ int main(int argc, char **argv) { cudaCheck(cudaSetDevice(deviceIdx)); // create host memory of random numbers - float* logits = make_random_float_01(B * T * V); - float* probs = (float*)malloc(B * T * V * sizeof(float)); + float* logits = make_random_float(B * T * V); + float* probs = make_random_float_01(B * T * V); float* dlogits = (float*)malloc(B * T * V * sizeof(float)); float* losses = (float*)malloc(B * T * sizeof(float)); float* dlosses = make_random_float(B * T); @@ -760,6 +760,7 @@ int main(int argc, char **argv) { free(losses); free(dlosses); free(targets); + free(outliers); cudaCheck(cudaFree(d_dlogits)); cudaCheck(cudaFree(d_losses)); cudaCheck(cudaFree(d_logits)); diff --git a/dev/cuda/crossentropy_softmax_backward.cu b/dev/cuda/crossentropy_softmax_backward.cu index 27521bf..65c72a2 100644 --- a/dev/cuda/crossentropy_softmax_backward.cu +++ b/dev/cuda/crossentropy_softmax_backward.cu @@ -99,7 +99,7 @@ int main(int argc, char **argv) { cudaCheck(cudaSetDevice(deviceIdx)); // create host memory of random numbers - float* probs = make_random_float(B * T * V); + float* probs = make_random_float_01(B * T * V); int* targets = make_random_int(B * T, V); float* dlosses = make_random_float(B * T); float* dlogits = make_zeros_float(B * T * V); diff --git a/llmc/gelu.cuh b/llmc/gelu.cuh index ce9aa12..cd5c297 100644 --- a/llmc/gelu.cuh +++ b/llmc/gelu.cuh @@ -50,7 +50,7 @@ __global__ void gelu_backward_inplace_kernel(floatX* d_in_out, const floatX* inp 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); + assert(N % (block_size * x128::size) == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); gelu_forward_kernel2<<>>(out, inp); cudaCheck(cudaGetLastError()); @@ -59,7 +59,7 @@ void gelu_forward(floatX* out, const floatX* inp, int N, cudaStream_t stream) { 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); + assert(N % (block_size * x128::size) == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); gelu_backward_inplace_kernel<<>>(d_in_out, inp); cudaCheck(cudaGetLastError()); diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index a49f0bf..9c9afa3 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -368,7 +368,7 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, 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); + assert(N % (block_size * x128::size) == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); residual_forward_kernel<<>>(out, inp1, inp2); cudaCheck(cudaGetLastError()); From 9b5ec8f719bb8471d397ad15fd394173ff5e499a Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 17:44:00 +0200 Subject: [PATCH 20/20] Set v memory to random floats >=0 --- dev/cuda/adamw.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dev/cuda/adamw.cu b/dev/cuda/adamw.cu index f27600a..74dfc2e 100644 --- a/dev/cuda/adamw.cu +++ b/dev/cuda/adamw.cu @@ -160,7 +160,7 @@ int main(int argc, char **argv) { float* params_memory = make_random_float(num_parameters); float* grads_memory = make_random_float(num_parameters); float* m_memory = make_random_float(num_parameters); - float* v_memory = make_random_float(num_parameters); + float* v_memory = make_random_float_01(num_parameters); // move to GPU float* d_params_memory;