Merge branch 'gradient-clipping' of https://github.com/ngc92/llm.c into ngc92-gradient-clipping

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

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 <assert.h>
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+// 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<class T>
+__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<float>{});
+    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<float>{});
+        if(warp.thread_rank() ==  0) {
+            atomicAdd(out, block_sum);
+        }
+    }
+}
+
+
+
+template<class T>
+__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<float>{});
+    // and atomic in global buffer
+    if(warp.thread_rank() == 0) {
+        atomicAdd(out, warp_result);
+    }
+}
+
+
+
+template<typename T>
+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<<<grid_size, block_size>>>(out, values, count);
+    cudaCheck(cudaGetLastError());
+}
+
+template<typename T>
+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<<<grid_size, block_size>>>(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));
+}

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

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.

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

train_gpt2.cu

@@ -1151,11 +1151,31 @@ __device__ float lerp(float start, float end, float weight) {
 template <typename Tp, typename Tg>
 __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<class T>
+__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<warpReduceSum>(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<typename T>
+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<<<grid_size, block_size>>>(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<<<num_blocks, block_size>>>(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);

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)