From d2f06aed015939827e0f6ffafc503bad4cf3ae52 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 1 Jul 2024 11:24:16 +0200 Subject: [PATCH 1/5] Fix periodic spikes - clean loss after val --- train_gpt2.cu | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 60e2450..33ef2f9 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -736,6 +736,7 @@ float gpt2_validate(GPT2 *model, const int* inputs, const int* targets, size_t B mean_loss += model->cpu_losses[i]; } mean_loss /= B*T; + cudaCheck(cudaMemset(acts.losses, 0, model->batch_size * model->seq_len * sizeof(float))); cudaCheck(cudaDeviceSynchronize()); return mean_loss; } @@ -1403,16 +1404,16 @@ void error_usage() { // main training loop int main(int argc, char *argv[]) { // read in the (optional) command line arguments - const char* train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; - const char* val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; - const char* load_filename = "gpt2_124M_bf16.bin"; // bf16 weights of the model + const char* train_data_pattern = "/hdd/llmc/fineweb/bin/fineweb_train_*.bin"; + const char* val_data_pattern = "/hdd/llmc/fineweb/bin/fineweb_val_*.bin"; + const char* load_filename = "d12"; // bf16 weights of the model const char* lr_scheduler_type = "cosine"; - const char* output_log_dir = NULL; + const char* output_log_dir = "/hdd/llmc/log_gpt2_124M/"; int checkpoint_every = 0; // write checkpoints every how many steps? int checkpoints_keep = 0; // how long checkpoint history do we keep? (in units of checkpoints) int major_checkpoint_every = 0; // major checkpoints never get deleted when maintaining history int resume = 0; // resume the optimization, if one is found inside output_log_dir? - int B = 4; // batch size + int B = 16; // batch size int T = 1024; // sequence length max int total_batch_size = -1; // will be calculated down below later, if not provided float learning_rate = 3e-4f; @@ -1421,15 +1422,15 @@ int main(int argc, char *argv[]) { float weight_decay = 0.0f; float skip_update_lossz = 0.0f; // skip update if loss goes above this in zscore float skip_update_gradz = 0.0f; // skip update if grad_norm goes above this in zscore - int val_loss_every = 20; // every how many steps do we eval validation loss? + int val_loss_every = 50; // every how many steps do we eval validation loss? int val_max_steps = 20; // how many batches max do we eval for validation loss? - int sample_every = 20; // every how many steps to do inference? + int sample_every = 25000; // every how many steps to do inference? int genT = 64; // number of steps of inference we will do int overfit_single_batch = 0; // useful for debugging, 1 = only load a single data batch once int max_steps = -1; int override_enable_tf32 = 1; int use_master_weights = 1; - int recompute = 1; // recompute during backward setting, 0 = none, 1 = recompute gelu + int recompute = 0; // recompute during backward setting, 0 = none, 1 = recompute gelu int zero_stage = 0; // Zero Optimization Stage for Multi-GPU training int hellaswag_eval = 0; // multi-node settings From de8454d2f487fb842fd1d82cae872c87e5d203e2 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 1 Jul 2024 11:30:36 +0200 Subject: [PATCH 2/5] Make the grad cleaning more robust --- train_gpt2.cu | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 33ef2f9..77e8bff 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -736,7 +736,6 @@ float gpt2_validate(GPT2 *model, const int* inputs, const int* targets, size_t B mean_loss += model->cpu_losses[i]; } mean_loss /= B*T; - cudaCheck(cudaMemset(acts.losses, 0, model->batch_size * model->seq_len * sizeof(float))); cudaCheck(cudaDeviceSynchronize()); return mean_loss; } @@ -755,14 +754,15 @@ void gpt2_zero_grad(GPT2 *model) { void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int grad_accum_steps, bool last_step) { NVTX_RANGE_FN(); + // init gradients of parameters and activations to zero + gpt2_zero_grad(model); + // lazily allocate the memory for gradients of the weights and activations, if needed if (model->grads_memory == NULL) { NvtxRange rng("InitGrads"); // allocate buffers for weight gradients printf0("allocating %d MiB for parameter gradients\n", (int)round(model->num_parameters * sizeof(floatX) / (1024 * 1024))); model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_elements, model->param_sizeof); - // init gradients of parameters and activations to zero - gpt2_zero_grad(model); // initialise cpu scratch buffers for encoder backward size_t num_c_groups = CEIL_DIV(model->config.channels, (WARP_SIZE * x128::size)); assert((size_t)(model->batch_size * model->seq_len) * num_c_groups < (1ULL<<31ULL)); // todo - maybe an issue for llama3-400B(?) @@ -1814,8 +1814,6 @@ int main(int argc, char *argv[]) { float grad_scale = (grad_norm > grad_clip) ? grad_clip / grad_norm : 1.0f; gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, grad_scale, step+1, &multi_gpu_config); } - // zero out the gradients for the next iteration - gpt2_zero_grad(&model); cudaCheck(cudaEventRecord(end)); cudaCheck(cudaEventSynchronize(end)); // wait for the end event to finish to get correct timings // --------------- TRAINING SECTION END ------------------- From 78e76278598ce3dc9e15ef3618397b0768376e24 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 1 Jul 2024 11:31:55 +0200 Subject: [PATCH 3/5] Remove local args for testing --- train_gpt2.cu | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 77e8bff..203e3aa 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1404,16 +1404,16 @@ void error_usage() { // main training loop int main(int argc, char *argv[]) { // read in the (optional) command line arguments - const char* train_data_pattern = "/hdd/llmc/fineweb/bin/fineweb_train_*.bin"; - const char* val_data_pattern = "/hdd/llmc/fineweb/bin/fineweb_val_*.bin"; - const char* load_filename = "d12"; // bf16 weights of the model + const char* train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; + const char* val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; + const char* load_filename = "gpt2_124M_bf16.bin"; // bf16 weights of the model const char* lr_scheduler_type = "cosine"; - const char* output_log_dir = "/hdd/llmc/log_gpt2_124M/"; + const char* output_log_dir = NULL; int checkpoint_every = 0; // write checkpoints every how many steps? int checkpoints_keep = 0; // how long checkpoint history do we keep? (in units of checkpoints) int major_checkpoint_every = 0; // major checkpoints never get deleted when maintaining history int resume = 0; // resume the optimization, if one is found inside output_log_dir? - int B = 16; // batch size + int B = 4; // batch size int T = 1024; // sequence length max int total_batch_size = -1; // will be calculated down below later, if not provided float learning_rate = 3e-4f; @@ -1422,15 +1422,15 @@ int main(int argc, char *argv[]) { float weight_decay = 0.0f; float skip_update_lossz = 0.0f; // skip update if loss goes above this in zscore float skip_update_gradz = 0.0f; // skip update if grad_norm goes above this in zscore - int val_loss_every = 50; // every how many steps do we eval validation loss? + int val_loss_every = 20; // every how many steps do we eval validation loss? int val_max_steps = 20; // how many batches max do we eval for validation loss? - int sample_every = 25000; // every how many steps to do inference? + int sample_every = 20; // every how many steps to do inference? int genT = 64; // number of steps of inference we will do int overfit_single_batch = 0; // useful for debugging, 1 = only load a single data batch once int max_steps = -1; int override_enable_tf32 = 1; int use_master_weights = 1; - int recompute = 0; // recompute during backward setting, 0 = none, 1 = recompute gelu + int recompute = 1; // recompute during backward setting, 0 = none, 1 = recompute gelu int zero_stage = 0; // Zero Optimization Stage for Multi-GPU training int hellaswag_eval = 0; // multi-node settings From 74e5ba72608aa02bdfcc79e1d60021ec0258d87f Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 1 Jul 2024 12:11:02 +0200 Subject: [PATCH 4/5] Add micro step guard --- profile_gpt2.cu | 2 +- test_gpt2.cu | 6 +++--- train_gpt2.cu | 10 ++++++---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/profile_gpt2.cu b/profile_gpt2.cu index 29fbdb9..f5efed8 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -61,7 +61,7 @@ int main(int argc, char *argv[]) { // do a training step gpt2_forward(&model, x, B, T); gpt2_zero_grad(&model); - gpt2_backward_and_reduce(&model, x, y, 1, true); + gpt2_backward_and_reduce(&model, x, y, 1, 0); float grad_norm = gpt2_calculate_grad_norm(&model, &multi_gpu_config); float grad_scale = (grad_norm > 1.0f) ? 1.0f / grad_norm : 1.0f; gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, grad_scale, 1, &multi_gpu_config); diff --git a/test_gpt2.cu b/test_gpt2.cu index 800df69..d7464bf 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -218,7 +218,7 @@ int main(int argc, char *argv[]) { clock_gettime(CLOCK_MONOTONIC, &start); gpt2_forward(&model, x, B, T); gpt2_zero_grad(&model); - gpt2_backward_and_reduce(&model, x, y, 1, true); + gpt2_backward_and_reduce(&model, x, y, 1, 0); clock_gettime(CLOCK_MONOTONIC, &end); double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9; @@ -336,7 +336,7 @@ int main(int argc, char *argv[]) { dataloader_next_batch(&loader); gpt2_forward(&model, loader.inputs, B, T); gpt2_zero_grad(&model); - gpt2_backward_and_reduce(&model, loader.inputs, loader.targets, 1, true); + gpt2_backward_and_reduce(&model, loader.inputs, loader.targets, 1, 0); gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+11, &multi_gpu_config); losses[step] = model.mean_loss; tokens[step] = loader.inputs[0]; @@ -351,7 +351,7 @@ int main(int argc, char *argv[]) { dataloader_next_batch(&loader); gpt2_forward(&model, loader.inputs, B, T); gpt2_zero_grad(&model); - gpt2_backward_and_reduce(&model, loader.inputs, loader.targets, 1, true); + gpt2_backward_and_reduce(&model, loader.inputs, loader.targets, 1, 0); gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+11, &multi_gpu_config); if(loader.inputs[0] != tokens[step]) { diff --git a/train_gpt2.cu b/train_gpt2.cu index 203e3aa..a2ccdcc 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -751,11 +751,13 @@ void gpt2_zero_grad(GPT2 *model) { cudaCheck(cudaDeviceSynchronize()); } -void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int grad_accum_steps, bool last_step) { +void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int grad_accum_steps, int micro_step) { NVTX_RANGE_FN(); + bool last_step = micro_step == grad_accum_steps - 1; - // init gradients of parameters and activations to zero - gpt2_zero_grad(model); + if (micro_step == 0) { + gpt2_zero_grad(model); // set correct grad state on first micro-step + } // lazily allocate the memory for gradients of the weights and activations, if needed if (model->grads_memory == NULL) { @@ -1795,7 +1797,7 @@ int main(int argc, char *argv[]) { // forward pass. note that we pass in grad_accum_steps, which scales down the loss gpt2_forward(&model, train_loader.inputs, B, T); // backward pass. all model params accumulate gradients with += inside this inner loop - gpt2_backward_and_reduce(&model, train_loader.inputs, train_loader.targets, grad_accum_steps, micro_step == grad_accum_steps - 1); + gpt2_backward_and_reduce(&model, train_loader.inputs, train_loader.targets, grad_accum_steps, micro_step); } float zloss = (float)(update_detector(&loss_outlier_detector, (double)model.mean_loss)); // loss z-score // fetch the next learning rate From 7b929300217ff1a974b63791a228928b39b26409 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Mon, 1 Jul 2024 17:20:46 +0000 Subject: [PATCH 5/5] move things around a bit, a few comments, and also update python script to correspond to C more --- train_gpt2.cu | 8 +++++--- train_gpt2.py | 6 +----- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index a56d181..246d413 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -740,10 +740,11 @@ float gpt2_validate(GPT2 *model, const int* inputs, const int* targets, size_t B return mean_loss; } - void gpt2_zero_grad(GPT2 *model) { NVTX_RANGE_FN(); - // the losses accumulate over the duration of gradient accumulation micro steps, also reset here + // there are currently two state vars during the gradient accumulation inner loop: + // 1) the losses accumulate += into acts.losses, reset here + // 2) the gradients accumulate += into grads_memory, reset here cudaCheck(cudaMemset(model->acts.losses, 0, model->batch_size * model->seq_len * sizeof(float))); if (model->grads_memory != NULL) { cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX))); @@ -755,8 +756,9 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int NVTX_RANGE_FN(); bool last_step = micro_step == grad_accum_steps - 1; + // on the first micro-step zero the gradients, as we're about to += accumulate into them if (micro_step == 0) { - gpt2_zero_grad(model); // set correct grad state on first micro-step + gpt2_zero_grad(model); } // lazily allocate the memory for gradients of the weights and activations, if needed diff --git a/train_gpt2.py b/train_gpt2.py index e5974f7..b9dee87 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -698,10 +698,6 @@ if __name__ == "__main__": write_state(model, x, y, logits, loss, f"gpt2_{model_size_str}_debug_state.bin") # reset the train_loader for the optimization below train_loader.reset() - # clear the grads here explicitly because otherwise we'd have a duplicate grad accumulation - # since in the training loop we do a backward() and then zero_grad() at the end of the loop - # this would cause an incorrect first training step - model.zero_grad() # ------------------------------------------------------------------------- # main training loop @@ -794,6 +790,7 @@ if __name__ == "__main__": # --------------- TRAINING SECTION BEGIN ----------------- model.train() + optimizer.zero_grad(set_to_none=True) # if we are trying to overfit a single batch, we reset the loader here if args.overfit_single_batch: train_loader.reset() @@ -830,7 +827,6 @@ if __name__ == "__main__": param_group['lr'] = lr # step the optimizer optimizer.step() - optimizer.zero_grad(set_to_none=True) # --------------- TRAINING SECTION END ------------------- # everything that follows now is just diagnostics, prints, logging, etc.