diff --git a/dev/cuda/Makefile b/dev/cuda/Makefile index c741788..14eae20 100644 --- a/dev/cuda/Makefile +++ b/dev/cuda/Makefile @@ -18,7 +18,7 @@ MPI_PATHS = -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux- $(NVCC) $(CFLAGS) $(NVCCFLAGS) $< -o $@ # Build all targets -TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward fused_residual_forward +TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward fused_residual_forward global_norm all: $(TARGETS) # Individual targets: forward pass @@ -48,6 +48,7 @@ matmul_backward: matmul_backward.cu # Update kernels adamw: adamw.cu +global_norm: global_norm.cu # NCCL communication kernels nccl_all_reduce: nccl_all_reduce.cu diff --git a/dev/cuda/global_norm.cu b/dev/cuda/global_norm.cu new file mode 100644 index 0000000..2295c49 --- /dev/null +++ b/dev/cuda/global_norm.cu @@ -0,0 +1,181 @@ +/* +Kernels for a global norm. +Global norm in this context means that we want to calculate a single norm cooperatively using all avalailable SMs, instead + of multiple norms that can be handled by separate blocks. + +Compile example: +nvcc -O3 --use_fast_math global_norm.cu -o global_norm +*/ + + +#include +#include +#include + +// turn on bf16 as default, done up here for now +#define ENABLE_BF16 +#include "common.h" + + +float global_norm_cpu(const float* data, size_t count) { + // accumulate in double so we have an accurate numerical reference + double acc = 0.0; + for(size_t i = 0; i < count; ++i) { + acc += (double)data[i] * (double)data[i]; + } + return (float)acc; +} + + +template +__global__ void norm_kernel1(float* out, 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. + namespace cg = cooperative_groups; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + + __shared__ float block_result[32]; + + // out will be updated atomically from all thread blocks + size_t index = threadIdx.x + blockDim.x * blockIdx.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]; + } + // warp-level reduce + float warp_result = cg::reduce(warp, accumulator, cg::plus{}); + block_result[warp.meta_group_rank()] = warp_result; + block.sync(); + if(warp.meta_group_rank() == 0) { + float gather = warp.thread_rank() < warp.meta_group_size() ? block_result[warp.thread_rank()] : 0.f; + float block_sum = cg::reduce(warp, gather, cg::plus{}); + if(warp.thread_rank() == 0) { + atomicAdd(out, block_sum); + } + } +} + + + +template +__global__ void norm_kernel2(float* out, const T* data, size_t count) { + // no shared memory; but one atomic per warp instead of per block + namespace cg = cooperative_groups; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + + // out will be updated atomically from all thread blocks + size_t index = threadIdx.x + blockDim.x * blockIdx.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]; + } + + // warp-level reduce + float warp_result = cg::reduce(warp, accumulator, cg::plus{}); + // and atomic in global buffer + if(warp.thread_rank() == 0) { + atomicAdd(out, warp_result); + } +} + + + +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. + // 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 = cuda_threads_per_SM * cuda_num_SMs / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + norm_kernel1<<>>(out, values, count); + cudaCheck(cudaGetLastError()); +} + +template +void global_norm2(float* out, const T* values, size_t count, int block_size) { + // ditto + const int grid_size = cuda_threads_per_SM * cuda_num_SMs / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + norm_kernel2<<>>(out, values, count); + cudaCheck(cudaGetLastError()); +} + +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); + } +} + +int main(int argc, const char **argv) { + setup_main(); + + int C = 768; + int L = 12; + + size_t num_params = (size_t)(C * 4*C + C*C) * 2 * L; + + // create host memory of random numbers + float* inp = make_random_float(num_params); + // scale them down + for(size_t i = 0; i < num_params; ++i) { + inp[i] *= 1e-3; + } + + // read kernel_num from command line + int kernel_num = 1; + if (argc > 1) { + kernel_num = atoi(argv[1]); + } + printf("Using kernel %d\n", kernel_num); + + // first check the correctness of the kernel + float out = global_norm_cpu(inp, num_params); + + // move to GPU + float* d_out; + floatX* d_inp; + cudaCheck(cudaMalloc(&d_out, sizeof(float))); + cudaCheck(cudaMalloc(&d_inp, num_params * sizeof(floatX))); + cudaCheck(memcpy_convert(d_inp, inp, num_params)); + + int block_sizes[] = {32, 64, 128, 256, 512, 768, 1024}; + for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { + int block_size = block_sizes[j]; + 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); + validate_result(d_out, &out, "out", 1, 1e-2f); + } + + printf("All results match. Starting benchmarks.\n\n"); + + for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) { + int block_size = block_sizes[j]; + + int repeat_times = 1000; + + float elapsed_time = benchmark_kernel(repeat_times, global_norm, + kernel_num, d_out, d_inp, + num_params, block_size); + size_t memory_ops = num_params * sizeof(floatX); + float memory_bandwidth = memory_ops / elapsed_time / 1e6; + + printf("block_size %4d | time %.4f ms | bandwidth %.2f GB/s\n", block_size, elapsed_time, memory_bandwidth); + } + + // free memory + free(inp); + cudaCheck(cudaFree(d_out)); + cudaCheck(cudaFree(d_inp)); +} \ No newline at end of file diff --git a/profile_gpt2.cu b/profile_gpt2.cu index 4b24c89..5a67645 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -54,7 +54,7 @@ int main(int argc, char *argv[]) { gpt2_forward(&model, x, y, B, T); gpt2_zero_grad(&model); gpt2_backward(&model); - gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, 1, &multi_gpu_config); + gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, 1.f, 1, &multi_gpu_config); cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings // free diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index d9dbd4f..de2edfd 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -130,7 +130,7 @@ for rid, row in kernel_profile_data: # the classifier part, counts only once pass_name = "cls" phase = "bwd" - elif "adamw" in kernel: + elif "adamw" in kernel or "global_norm" in kernel: # encoder layer or adam pass_name = "opt" # before the first optimizer run, we create weight copies. diff --git a/test_gpt2.cu b/test_gpt2.cu index 84701b0..50a291f 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -274,7 +274,7 @@ int main(int argc, char *argv[]) { allok = allok & check_tensor(tensors1[15], tensors2[15], C, "lnfb", 2e-2f); } - gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.01f, step+1, &multi_gpu_config); + gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.01f, 1.f, step+1, &multi_gpu_config); // print the timing information at the end printf("step %d: loss %f (took %f ms)\n", step+1, model.mean_loss, time_elapsed_s * 1000); @@ -283,16 +283,16 @@ int main(int argc, char *argv[]) { // expected losses are as follows, from Python float expected_losses[10] = { - 5.270007133483887, - 4.059706687927246, - 3.3751230239868164, - 2.8007826805114746, - 2.315382242202759, - 1.8490285873413086, - 1.3946564197540283, - 0.9991465210914612, - 0.6240804195404053, - 0.37651097774505615 + 5.2700, + 4.0607, + 3.3166, + 2.7115, + 2.1702, + 1.6349, + 1.1419, + 0.7038, + 0.3769, + 0.1743 }; // compare diff --git a/train_gpt2.cu b/train_gpt2.cu index 030c5c9..05a0ba2 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1151,11 +1151,31 @@ __device__ float lerp(float start, float end, float weight) { template __global__ void adamw_kernel3(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, float beta1_correction, float beta2_correction, float eps, float weight_decay, + float* grad_norm, float max_grad_norm, unsigned int seed) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_parameters) { return; } // guard + + float scale = 1.f; + if(!isfinite(*grad_norm)) { + // if we had a numerical problem (e.g, overflow) + // in our gradient calculation, don't mess up the + // existing weights. + // TODO increase a global counter somewhere so we actually know if/how often this happens + if(threadIdx.x == 0 && blockIdx.x == 0) { + printf("[WARNING] weight update skipped due to non-finite gradients!\n"); + } + return; + } + if(*grad_norm > max_grad_norm) { + scale = max_grad_norm / sqrtf(*grad_norm); + // TODO just for debugging, remove this + if(threadIdx.x == 0 && blockIdx.x == 0) { + printf("[norm %f]\n", sqrtf(*grad_norm)); + } + } // get the gradient, m, and v for this parameter - float grad = (float)grads_memory[idx]; + float grad = scale * (float)grads_memory[idx]; float m = m_memory[idx]; float v = v_memory[idx]; // update the first moment (momentum) @@ -1180,6 +1200,27 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg if (master_params_memory != NULL) { master_params_memory[idx] = param; } } +template +__global__ void global_norm_kernel(float* out, 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 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]; + } + // warp-level reduce + float block_sum = blockReduce(accumulator); + if(threadIdx.x == 0) { + atomicAdd(out, block_sum); + } +} + struct SoftmaxParams { float Scale; float Offset; @@ -1656,6 +1697,20 @@ void fused_classifier(Type* logits, Type* losses, cudaCheck(cudaGetLastError()); } +template +void global_norm(float* out, const T* values, size_t count) { + 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 + // 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 + global_norm_kernel<<>>(out, values, count); + cudaCheck(cudaGetLastError()); +} + // ---------------------------------------------------------------------------- // GPT-2 model definition @@ -2333,7 +2388,7 @@ void gpt2_multi_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) { #endif } -void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t, MultiGpuConfig* multi_gpu_config) { +void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, float grad_clipping, int t, MultiGpuConfig* multi_gpu_config) { NVTX_RANGE_FN(); size_t num_parameters = multi_gpu_config->shard_num_parameters; floatX* params_memory = (floatX*)model->params_memory + multi_gpu_config->shard_offset; @@ -2354,6 +2409,14 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo } } + // repurposing this buffer. We calculate the gradient norm on the GPU, and need it in the next kernel, + // so we _really_ don't want to transfer it here as an actual float. So we just pass around a pointer + // to this memory that is not otherwise needed during the update phase. + float* grad_norm = (float*)model->acts.output; + + // global gradient norm + global_norm(grad_norm, (floatX*)model->grads_memory, model->num_parameters); + int block_size = 512; int num_blocks = CEIL_DIV(num_parameters, block_size); float beta1_correction = 1.0f - powf(beta1, t); @@ -2361,7 +2424,9 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo unsigned int seed = random_u32(&model->rng_state); adamw_kernel3<<>>(params_memory, model->master_weights, grads_memory, model->m_memory, model->v_memory, num_parameters, - learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, seed); + learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, + grad_norm, grad_clipping, + seed); cudaCheck(cudaGetLastError()); } @@ -2607,6 +2672,7 @@ int main(int argc, char *argv[]) { int use_master_weights = 1; int recompute = 1; // recompute during backward setting, 0 = none, 1 = recompute gelu int zero_stage = 0; // Zero Optimization Stage for Multi-GPU training + float grad_clipping = 1.f; for (int i = 1; i < argc; i+=2) { if (i + 1 >= argc) { error_usage(); } // must have arg after flag if (argv[i][0] != '-') { error_usage(); } // must start with dash @@ -2627,6 +2693,7 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 'a') { overfit_single_batch = atoi(argv[i+1]); } else if (argv[i][1] == 'f') { override_enable_tf32 = atoi(argv[i+1]); } else if (argv[i][1] == 'w') { use_master_weights = atoi(argv[i+1]); } + else if (argv[i][1] == 'c') { grad_clipping = atof(argv[i+1]); } else if (argv[i][1] == 'z') { zero_stage = atoi(argv[i+1]); } else if (argv[i][1] == 'r') { recompute = atoi(argv[i+1]); } else { error_usage(); } @@ -2642,6 +2709,7 @@ int main(int argc, char *argv[]) { printf0("| sequence length T | %-50d |\n", T); printf0("| total batch size | %-50d |\n", total_batch_size); printf0("| learning rate | %-50e |\n", learning_rate); + printf0("| grad_clipping | %-50e |\n", grad_clipping); printf0("| max_steps | %-50d |\n", max_steps); printf0("| val_loss_every | %-50d |\n", val_loss_every); printf0("| val_max_batches | %-50d |\n", val_max_batches); @@ -2826,7 +2894,7 @@ int main(int argc, char *argv[]) { model.mean_loss = lossf; // update the parameters gpt2_multi_gpu_accumulate(&model, &multi_gpu_config); - gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1, &multi_gpu_config); + gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, grad_clipping, step+1, &multi_gpu_config); gpt2_multi_gpu_gather(&model, &multi_gpu_config); // zero out the gradients for the next iteration gpt2_zero_grad(&model); diff --git a/train_gpt2.py b/train_gpt2.py index c50fc7b..b57fe43 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -407,6 +407,7 @@ if __name__ == "__main__": parser.add_argument("--batch_size", type=int, default=4, help="batch size, in units of #batch dimensions") parser.add_argument("--sequence_length", type=int, default=64, help="sequence length") parser.add_argument("--total_batch_size", type=int, default=256, help="total desired batch size, in units of #tokens") + parser.add_argument("--grad_clipping", type=float, default=1, help="maximum gradient magnitude") args = parser.parse_args() B, T = args.batch_size, args.sequence_length assert 1 <= T <= 1024 @@ -552,6 +553,7 @@ if __name__ == "__main__": if device == "cuda": torch.cuda.reset_peak_memory_stats() timings = [] + norm = -1 # dummy value to print in inference-only mode for step in range(args.num_iterations): t0 = time.time() @@ -575,7 +577,7 @@ if __name__ == "__main__": # backward pass if not args.inference_only: loss.backward() - # todo: grad clip here + norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clipping) optimizer.step() optimizer.zero_grad(set_to_none=True) @@ -588,7 +590,7 @@ if __name__ == "__main__": t1 = time.time() # the 0th iteration is often an outlier (much slower) => skip logging it tokens_per_second = grad_accum_steps * ddp_world_size * B * T / (t1-t0) - print0(f"iteration {step+1}, loss: {lossf:.4f}, time: {(t1-t0)*1000:.3f}ms, tok/s: {tokens_per_second:.2f}") + print0(f"iteration {step+1}, loss: {lossf:.4f}, time: {(t1-t0)*1000:.3f}ms, tok/s: {tokens_per_second:.2f}, norm: {norm:.3f}") if step > 0 and step > args.num_iterations - 20: timings.append(t1-t0)