diff --git a/README.md b/README.md index f77870b..7b9c2d4 100644 --- a/README.md +++ b/README.md @@ -403,6 +403,9 @@ Lastly, I will be a lot more sensitive to complexity in the root folder of the p - [llm.rs](https://github.com/yijunyu/llm.rs) by @[Yijun Yu](https://github.com/yijunyu): a Rust rewrite with the aim to have same performance - [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project +- Swift + - [llm.swift](https://github.com/otabuzzman/llm.swift) by @[otabuzzman](https://github.com/otabuzzman): a Swift port of this project + - Zig - [llm.zig](https://github.com/Saimirbaci/llm.zig) by @[saimirbaci](https://github.com/Saimirbaci): a Zig port of this project diff --git a/rand.h b/rand.h new file mode 100644 index 0000000..e60e5e6 --- /dev/null +++ b/rand.h @@ -0,0 +1,223 @@ +/* +Mersenne Twisters implementation, numerically identical to torch. + +Example usage: + + mt19937_state state; + manual_seed(&state, 137); + printf("%u\n", randint32(&state)); + printf("%u\n", randint32(&state)); + printf("%u\n", randint32(&state)); + printf("%u\n", randint32(&state)); + printf("%u\n", randint32(&state)); + + float t8[8]; + normal_(t8, 8, 0, 1, &state); + for (int i = 0; i < 8; i++) { + printf("%f\n", t8[i]); + } + printf("%u\n", randint32(&state)); + + float t16[16]; + normal_(t16, 16, 0, 1, &state); + for (int i = 0; i < 16; i++) { + printf("%f\n", t16[i]); + } + printf("%u\n", randint32(&state)); + +PyTorch reference (producing identical results): + + import torch + torch.manual_seed(137) + print(torch.randint(0, 0xFFFFFFFF, [1]).item()) + print(torch.randint(0, 0xFFFFFFFF, [1]).item()) + print(torch.randint(0, 0xFFFFFFFF, [1]).item()) + print(torch.randint(0, 0xFFFFFFFF, [1]).item()) + print(torch.randint(0, 0xFFFFFFFF, [1]).item()) + t = torch.zeros(8); + t.normal_() + for i in range(len(t)) : + print(t[i].item()) + print(torch.randint(0, 0xFFFFFFFF, [1]).item()) + t = torch.zeros(16); + t.normal_() + for i in range(len(t)) : + print(t[i].item()) + print(torch.randint(0, 0xFFFFFFFF, [1]).item()) + +Both output: + + 4053805790 + 2173880614 + 380293709 + 1237255315 + 2986595568 + 0.7947664260864258 + 1.4369317293167114 + - 0.2292192131280899 + 0.47556325793266296 + - 0.6334410905838013 + - 0.5791953802108765 + - 0.0925704762339592 + - 0.8659197092056274 + 2186503452 + - 1.2813878059387207 + - 2.646395683288574 + - 0.06569503247737885 + 0.2180829495191574 + - 0.46536165475845337 + - 0.33108410239219666 + 2.5485482215881348 + 0.10425379872322083 + 0.8460659980773926 + 0.9462448358535767 + - 0.2913765013217926 + 0.34313806891441345 + - 1.1186704635620117 + - 0.18305328488349915 + - 2.3153159618377686 + 0.3961987793445587 + 2756748748 +*/ + +#ifndef RAND_H +#define RAND_H + +#include + +#define MERSENNE_STATE_M 397u +#define MERSENNE_STATE_N 624u + +#define LMASK 0x7ffffffful +#define UMASK 0x80000000ul + +// Copyright(c) Makoto Matsumoto and Takuji Nishimura + +// This implementation follows PyTorch so that we are numerically identical when running verification tests. + +typedef struct { + unsigned long long seed_; + int left_; + unsigned int next_; + unsigned int state_[MERSENNE_STATE_N]; + unsigned int MATRIX_A[2]; +} mt19937_state; + +void manual_seed(mt19937_state* state, unsigned int seed) { + state->MATRIX_A[0] = 0x0u; + state->MATRIX_A[1] = 0x9908b0df; + state->state_[0] = seed & 0xffffffff; + for (unsigned int j = 1; j < MERSENNE_STATE_N; j++) { + state->state_[j] = 1812433253 * (state->state_[j - 1] ^ (state->state_[j - 1] >> 30)) + j; + state->state_[j] &= 0xffffffff; + } + state->left_ = 1; + state->next_ = 0; +} + +void next_state(mt19937_state* state) { + state->left_ = MERSENNE_STATE_N; + state->next_ = 0; + unsigned int y, j; + for (j = 0; j < MERSENNE_STATE_N - MERSENNE_STATE_M; j++) { + y = (state->state_[j] & UMASK) | (state->state_[j + 1] & LMASK); + state->state_[j] = state->state_[j + MERSENNE_STATE_M] ^ (y >> 1) ^ state->MATRIX_A[y & 0x1]; + } + for (; j < MERSENNE_STATE_N - 1; j++) { + y = (state->state_[j] & UMASK) | (state->state_[j + 1] & LMASK); + state->state_[j] = state->state_[j + (MERSENNE_STATE_M - MERSENNE_STATE_N)] ^ (y >> 1) ^ state->MATRIX_A[y & 0x1]; + } + y = (state->state_[MERSENNE_STATE_N - 1] & UMASK) | (state->state_[0] & LMASK); + state->state_[MERSENNE_STATE_N - 1] = state->state_[MERSENNE_STATE_M - 1] ^ (y >> 1) ^ state->MATRIX_A[y & 0x1]; +} + +unsigned int randint32(mt19937_state* state) { + if (!state) return 0; + if (state->MATRIX_A[0] != 0 || state->MATRIX_A[1] != 0x9908b0df) manual_seed(state, 5489); // auto-initialize + if (--state->left_ <= 0) { + next_state(state); + } + unsigned int y = state->state_[state->next_++]; + y ^= y >> 11; + y ^= (y << 7) & 0x9d2c5680; + y ^= (y << 15) & 0xefc60000; + y ^= y >> 18; + return y; +} + +inline unsigned long long randint64(mt19937_state* state) { + return (((unsigned long long)(randint32(state)) << 32) | randint32(state)); +} + +inline float randfloat32(mt19937_state* state) { + return (randint32(state) & ((1ull << 24) - 1)) * (1.0f / (1ull << 24)); +} + +inline double randfloat64(mt19937_state* state) { + return (randint64(state) & ((1ull << 53) - 1)) * (1.0 / (1ull << 53)); +} + +void uniform_(float* data, unsigned int numel, float from, float to, mt19937_state* state) { + for (unsigned int t = 0; t < numel; t++) { + data[t] = randfloat32(state) * (to - from) + from; + } +} + +// Box�Muller transform + +void normal_fill_16(float* data, float mean, float std, mt19937_state* state) { + #define EPSILONE 1e-12 + for (unsigned int t = 0; t < 8; t++) { + float u1 = 1 - data[t]; + float u2 = data[t + 8]; + float radius = sqrtf(-2 * logf(u1 + EPSILONE)); + float theta = 2.0 * M_PI * u2; + data[t] = (radius * cosf(theta) * std + mean); + data[t + 8] = (radius * sinf(theta) * std + mean); + } +} + +void normal_fill(float* data, unsigned int numel, float mean, float std, mt19937_state* state) { + for (unsigned int t = 0; t < numel; t++) { + data[t] = randfloat32(state); + } + for (unsigned int i = 0; i < numel - 15; i += 16) { + normal_fill_16(data + i, mean, std, state); + } + if (numel % 16 != 0) { + // recompute the last 16 values + data = data + numel - 16; + for (unsigned int i = 0; i < 16; i++) { + data[i] = randfloat32(state); + } + normal_fill_16(data, mean, std, state); + } +} + +void normal_(float* data, unsigned int numel, float mean, float std, mt19937_state* state) { + #define EPSILONE 1e-12 + if (numel >= 16) { + normal_fill(data, numel, mean, std, state); + } + else { + double next_double_normal_sample; + int has_next_double_normal_sample = 0; + for (unsigned int t = 0; t < numel; t++) { + if (has_next_double_normal_sample) { + data[t] = (float)(next_double_normal_sample * std + mean); + has_next_double_normal_sample = 0; + continue; + } + // for numel < 16 we draw a double (float64) + float u1 = randfloat64(state); + float u2 = randfloat64(state); + float radius = sqrtf(-2 * logf(1 - u2 + EPSILONE)); + float theta = 2.0 * M_PI * u1; + next_double_normal_sample = radius * sinf(theta); + has_next_double_normal_sample = 1; + data[t] = (radius * cosf(theta) * std + mean); + } + } +} + +#endif \ No newline at end of file diff --git a/test_gpt2.cu b/test_gpt2.cu index 50a291f..247cd32 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, 1.f, step+1, &multi_gpu_config); + gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, 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); @@ -285,14 +285,14 @@ int main(int argc, char *argv[]) { float expected_losses[10] = { 5.2700, 4.0607, - 3.3166, - 2.7115, - 2.1702, - 1.6349, - 1.1419, - 0.7038, - 0.3769, - 0.1743 + 3.3202, + 2.7176, + 2.1811, + 1.6538, + 1.1680, + 0.7367, + 0.4008, + 0.1874 }; // compare diff --git a/train_gpt2.cu b/train_gpt2.cu index 798fe76..81f82f7 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -57,7 +57,9 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200), // defines: dataloader_init, dataloader_reset, dataloader_next_batch, dataloader_free // defines: evalloader_init, evalloader_reset, evalloader_next_batch, evalloader_free #include "dataloader.h" - +// defines: manual_seed, normal_ +// numerically identical to PyTorch's torch.manual_seed and torch.normal +#include "rand.h" // ---------------------------------------------------------------------------- // CUDA precision settings @@ -2035,6 +2037,110 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { model->recompute = 1; // default to recompute gelu during backward } +void gpt2_build_from_random(GPT2 *model, int depth) { + // init random (training from scratch) + + // parameterize the size of gpt2 based only on the depth of the model (num_layers) + model->config.num_layers = depth; + // follows GPT-2 sizes + int channels, num_heads; + if (depth == 12) { channels = 768; num_heads = 12; } // gpt2 (124M) + else if (depth == 24) { channels = 1024; num_heads = 16; } // gpt2-medium (350M) + else if (depth == 36) { channels = 1280; num_heads = 20; } // gpt2-large (774M) + else if (depth == 48) { channels = 1600; num_heads = 25; } // gpt2-xl (1558M) + else { fprintf(stderr, "Unsupported depth for now\n"); exit(EXIT_FAILURE); } + model->config.channels = channels; + model->config.num_heads = num_heads; + model->config.max_seq_len = 1024; + model->config.vocab_size = 50257; + model->config.padded_vocab_size = 50304; // padded to 128 + + // fill in all the parameter tensor dimensions and types + fill_in_parameter_sizes(model->param_elements, model->param_sizeof, model->config); + model->num_parameters = 0; + model->num_parameters_bytes = 0; + for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + model->num_parameters += model->param_elements[i]; + model->num_parameters_bytes += model->param_elements[i] * model->param_sizeof[i]; + } + // create memory for model parameters on the device + model->params_memory = malloc_and_point_parameters(&model->params, model->param_elements, model->param_sizeof); + + // allocate and random init the memory for all the parameters with GPT-2 schema + // weights ~N(0, 0.02), biases 0, c_proj weights ~N(0, 0.02/(2*L)**0.5) + // NOTE: assuming all parameters are of the type floatX, could be relaxed later + mt19937_state init_rng; + manual_seed(&init_rng, 42); + floatX* params_memory_cpu = (floatX*)mallocCheck(model->num_parameters_bytes); + memset(params_memory_cpu, 0, model->num_parameters_bytes); + // fill in all the weights with random values + float residual_scale = 1.0f / sqrtf(2.0f * model->config.num_layers); + // we have to init all these tensors exactly in the order that PyTorch initializes them + // so that we can match them up and get correctness and exactly the same initial conditions + size_t L = model->config.num_layers; + size_t offset = 0; + for (int l = 0; l < L; l++) { + offset = 0; + for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + // the layernorm parameters are all initialized to 1 + if (l == 0 && (i == 2 || i == 8 || i == 14)) { // only at l = 0 to init these just once + for (size_t j = 0; j < model->param_elements[i]; j++) { + params_memory_cpu[offset + j] = 1.0f; + } + } + // weights tensors are handled here + if ((l == 0 && (i == 0 || i == 1)) // only at l = 0, init the wte and wpe tensors + || i == 4 || i == 6 || i == 10 || i == 12) { + int n = model->param_elements[i]; + size_t layer_offset = 0; + if (i == 0) { + // for wte tensor (padded vocab) override to init V instead of Vp rows + n = model->config.vocab_size * model->config.channels; + } + if (i == 4 || i == 6 || i == 10 || i == 12) { + // weight tensors, we are only initializing layer l + assert(n % L == 0); + n = n / L; + layer_offset = l * n; + } + // in GPT-2, the projections back into the residual stream are additionally + // scaled by 1/sqrt(2*L) for training stability + float scale = (i == 6 || i == 12) ? 0.02f * residual_scale : 0.02f; + // okay let's draw the random numbers and write them + float *fp32_buffer = (float*)mallocCheck(n * sizeof(float)); + normal_(fp32_buffer, n, 0.0f, scale, &init_rng); + for (size_t j = 0; j < n; j++) { + params_memory_cpu[offset + layer_offset + j] = (floatX)fp32_buffer[j]; + } + free(fp32_buffer); + } + offset += model->param_elements[i]; + } + } + + // copy them to GPU + cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, model->num_parameters_bytes, cudaMemcpyHostToDevice)); + free(params_memory_cpu); + + // other inits and defaults + model->acts_memory = NULL; + model->grads_memory = NULL; + model->m_memory = NULL; + model->v_memory = NULL; + model->master_weights = NULL; + model->grads_acts_memory = NULL; + model->inputs = NULL; + model->targets = NULL; + model->cpu_losses = NULL; + model->cpu_losses_fp32 = NULL; + model->batch_size = 0; + model->seq_len = 0; + model->mean_loss = -1.0f; // -1.0f designates no loss + model->rng_state = 13371337; + model->use_master_weights = 1; // keep master weights copy in float for optim update? + model->recompute = 1; // default to recompute gelu during backward +} + void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, int grad_accum_steps=1) { NVTX_RANGE_FN(); // targets are optional and could be NULL @@ -2524,14 +2630,17 @@ typedef struct { } Logger; void logger_init(Logger *logger, const char *filename) { - logger->flush_every = 20; + logger->flush_every = 10; logger->logfile = NULL; - if (filename != NULL) { logger->logfile = fopenCheck(filename, "w"); } + // only rank 0 process will log + if (filename != NULL && multi_gpu_config.process_rank == 0) { + logger->logfile = fopenCheck(filename, "w"); + } } -void logger_log_eval(Logger *logger, int step, float val_loss) { +void logger_log_eval(Logger *logger, int step, float val) { if (logger->logfile != NULL) { - fprintf(logger->logfile, "s:%d eval:%.4f\n", step, val_loss); + fprintf(logger->logfile, "s:%d eval:%.4f\n", step, val); } } @@ -2544,7 +2653,7 @@ void logger_log_val(Logger *logger, int step, float val_loss) { void logger_log_train(Logger *logger, int step, float train_loss) { if (logger->logfile != NULL) { fprintf(logger->logfile, "s:%d trl:%.4f\n", step, train_loss); - if (step % 10 == 0) { fflush(logger->logfile); } + if (step % logger->flush_every == 0) { fflush(logger->logfile); } } } @@ -2566,6 +2675,9 @@ void error_usage() { fprintf(stderr, " -t sequence length T (default = 1024)\n"); fprintf(stderr, " -d total desired batch size (default = B * T * num_processes, i.e. no grad accumulation\n"); fprintf(stderr, " -l learning rate (default = 3e-4f)\n"); + fprintf(stderr, " -u learning rate warmup iterations (default = 0, no warmup)\n"); + fprintf(stderr, " -q learning rate decay: final fraction, at end of training (default = 1.0 (no decay))\n"); + fprintf(stderr, " -c weight decay (default = 0.0f)\n"); fprintf(stderr, " -x max_steps of optimization to run (-1 (default) = disable, run 1 epoch)\n"); fprintf(stderr, " -v val_loss_every, how often we evaluate val loss (default = 20)\n"); fprintf(stderr, " -m val_max_batches, up to how many val batches to estimate val loss? (default = 20)\n"); @@ -2576,6 +2688,7 @@ void error_usage() { fprintf(stderr, " -w keep f32 copy of weights for the optimizer? (default: 1)\n"); fprintf(stderr, " -z zero_stage, Zero Optimization Stage, 0,1,2,3 (default = 0)\n"); fprintf(stderr, " -r recompute: saves memory at cost of speed. (default = 1), 0 = none. 1 = recompute gelu\n"); + fprintf(stderr, " -h hellaswag eval run? (default = 0)\n"); exit(EXIT_FAILURE); } @@ -2593,6 +2706,9 @@ int main(int argc, char *argv[]) { 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; + int warmup_iterations = 0; + float final_learning_rate_frac = 1.0f; // final fraction of learning rate, at end of training + float weight_decay = 0.0f; int val_loss_every = 20; // every how many steps do we eval validation loss? int val_max_batches = 20; // how many batches max do we eval for validation loss? int sample_every = 20; // every how many steps to do inference? @@ -2604,6 +2720,7 @@ int main(int argc, char *argv[]) { 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_clip = 1.0f; + int hellaswag_eval = 0; 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 @@ -2617,6 +2734,9 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 't') { T = atoi(argv[i+1]); } else if (argv[i][1] == 'd') { total_batch_size = atoi(argv[i+1]); } else if (argv[i][1] == 'l') { learning_rate = atof(argv[i+1]); } + else if (argv[i][1] == 'u') { warmup_iterations = atoi(argv[i+1]); } + else if (argv[i][1] == 'q') { final_learning_rate_frac = atof(argv[i+1]); } + else if (argv[i][1] == 'c') { weight_decay = atof(argv[i+1]); } else if (argv[i][1] == 'x') { max_steps = atoi(argv[i+1]); } else if (argv[i][1] == 'v') { val_loss_every = atoi(argv[i+1]); } else if (argv[i][1] == 'm') { val_max_batches = atoi(argv[i+1]); } @@ -2628,8 +2748,11 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 'c') { grad_clip = 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 if (argv[i][1] == 'h') { hellaswag_eval = atoi(argv[i+1]); } else { error_usage(); } } + // should do a bit more error checking here + assert(warmup_iterations >= 0); // calculate a sensible default for total batch size by assuming no gradient accumulation if (total_batch_size == -1) { total_batch_size = B * T * multi_gpu_config.num_processes; } // if we're only overfitting a single batch for debugging, let's overfit the first batch @@ -2644,7 +2767,10 @@ int main(int argc, char *argv[]) { printf0("| micro batch size B | %-50d |\n", B); printf0("| sequence length T | %-50d |\n", T); printf0("| total batch size | %-50d |\n", total_batch_size); - printf0("| learning rate | %-50e |\n", learning_rate); + printf0("| learning rate (LR) | %-50e |\n", learning_rate); + printf0("| warmup iterations | %-50d |\n", warmup_iterations); + printf0("| final LR fraction | %-50e |\n", final_learning_rate_frac); + printf0("| weight decay | %-50e |\n", weight_decay); printf0("| grad_clip | %-50e |\n", grad_clip); printf0("| max_steps | %-50d |\n", max_steps); printf0("| val_loss_every | %-50d |\n", val_loss_every); @@ -2666,9 +2792,23 @@ int main(int argc, char *argv[]) { printf0("| precision | %-50s |\n", precision_str); printf0("+-----------------------+----------------------------------------------------+\n"); - // build the GPT-2 model from a checkpoint + // build the GPT-2 model GPT2 model; - gpt2_build_from_checkpoint(&model, load_filename); + // if load_filename is of the form "dX" where X is an integer (e.g. d12), then we build + // a random model with the depth of the model specified by X (e.g. 12). otherwise interpret + // this variable as a checkpoint filename, and load that checkpoint + assert(strlen(load_filename) >= 2); + if (load_filename[0] == 'd') { + int depth = atoi(load_filename + 1); + if (depth > 1 && depth <= 1000) { // we're not going to train models this big right? heh + gpt2_build_from_random(&model, depth); + } else { + exit(EXIT_FAILURE); + } + } else { + gpt2_build_from_checkpoint(&model, load_filename); + } + model.use_master_weights = use_master_weights; model.recompute = recompute; printf0("| load_filename | %-50s |\n", load_filename); @@ -2695,10 +2835,11 @@ int main(int argc, char *argv[]) { EvalLoader eval_loader; const char* hellaswag_path = "dev/data/hellaswag/hellaswag_val.bin"; const char hellaswag_available = access(hellaswag_path, F_OK) == 0; - if (hellaswag_available) { + const char run_hellaswag = hellaswag_eval && hellaswag_available; + if (run_hellaswag) { evalloader_init(&eval_loader, hellaswag_path, B, T, multi_gpu_config.process_rank, multi_gpu_config.num_processes); } - printf0("| hellaswag available | %-50s |\n", hellaswag_available ? "yes" : "no"); + printf0("| run hellaswag | %-50s |\n", run_hellaswag ? "yes" : "no"); printf0("+-----------------------+----------------------------------------------------+\n"); // pretty print in a table the multi-gpu configuration as well @@ -2710,7 +2851,7 @@ int main(int argc, char *argv[]) { // prints outside of pretty table to here and below if (!hellaswag_available) { printf0("HellaSwag eval not found at %s, skipping its evaluation\n", hellaswag_path); - printf0("You can run `python dev/data/hellaswag.py` to export and use it.\n"); + printf0("You can run `python dev/data/hellaswag.py` to export and use it with `-h 1`.\n"); } // more prints related to allocations from gpt2_build_from_checkpoint down here to not mess up our table above printf0("num_parameters: %zu => bytes: %zu\n", model.num_parameters, model.num_parameters_bytes); @@ -2724,9 +2865,11 @@ int main(int argc, char *argv[]) { B, T, multi_gpu_config.num_processes, total_batch_size); printf0("=> setting grad_accum_steps=%d\n", grad_accum_steps); - // set up the Logger & Tokenizer + // set up the Logger Logger logger; logger_init(&logger, output_log_file); + + // set up the Tokenizer Tokenizer tokenizer; tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); @@ -2765,8 +2908,8 @@ int main(int argc, char *argv[]) { } // once in a while estimate HellaSwag accuracy - if (hellaswag_available && - (step % val_loss_every == 0 || last_step)) { + if (run_hellaswag && + ((step > 0 && step % val_loss_every == 0) || last_step)) { NvtxRange evaluation_range("evaluation"); float eval_acc_norm = 0.0f; evalloader_reset(&eval_loader); @@ -2780,7 +2923,7 @@ int main(int argc, char *argv[]) { // careful because not all ranks may have the exact same allocation of number of examples eval_acc_norm = multi_gpu_cpu_float_sum(eval_acc_norm); printf0("HellaSwag: %d/%d = %f\n", (int)eval_acc_norm, eval_loader.num_examples, eval_acc_norm / eval_loader.num_examples); - logger_log_eval(&logger, step, eval_acc_norm); + logger_log_eval(&logger, step, eval_acc_norm / eval_loader.num_examples); } // once in a while do model inference to print generated text @@ -2857,7 +3000,20 @@ int main(int argc, char *argv[]) { model.mean_loss = lossf; // update the parameters gpt2_multi_gpu_accumulate(&model, &multi_gpu_config); - float grad_norm = gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, grad_clip, step+1, &multi_gpu_config); + // learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac + float step_learning_rate = learning_rate; + if (step < warmup_iterations) { + step_learning_rate = learning_rate * ((float)(step + 1)) / warmup_iterations; + } else { + float decay_ratio = ((float)(step - warmup_iterations)) / (train_num_batches - warmup_iterations); + assert(0.0f <= decay_ratio && decay_ratio <= 1.0f); + float coeff = 0.5f * (1.0f + cosf(M_PI * decay_ratio)); // coeff starts at 1 and goes to 0 + assert(0.0f <= coeff && coeff <= 1.0f); + float min_lr = learning_rate * final_learning_rate_frac; + step_learning_rate = min_lr + coeff * (learning_rate - min_lr); + } + // update the model parameters + float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, grad_clip, 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); @@ -2870,17 +3026,17 @@ int main(int argc, char *argv[]) { float time_elapsed_ms; cudaCheck(cudaEventElapsedTime(&time_elapsed_ms, start, end)); size_t tokens_processed = (size_t)multi_gpu_config.num_processes * B * T * grad_accum_steps; - float tokens_per_second = tokens_processed / time_elapsed_ms * 1000.0; + float tokens_per_second = tokens_processed / time_elapsed_ms * 1000.0f; float bias_corrected_ema_tokens_per_second = tokens_per_second; // by default set to non-ema version if (step > 0) { // consider the first batch to be a warmup (e.g. cuBLAS/cuDNN initialisation) - total_sum_iteration_time_s += time_elapsed_ms / 1000.0; + total_sum_iteration_time_s += time_elapsed_ms / 1000.0f; // smooth out the tok/s with an exponential moving average, and bias correct just like in AdamW ema_tokens_per_second = 0.95f * ema_tokens_per_second + 0.05f * tokens_per_second; bias_corrected_ema_tokens_per_second = ema_tokens_per_second / (1.0f - powf(0.95f, step)); } float accumulated_loss = multi_gpu_config.num_processes == 1 ? model.mean_loss : model.accumulated_mean_loss; - printf0("step %4d/%d: train loss %f norm %.4f (%.2f ms, %.0f tok/s)\n", - step + 1, train_num_batches, accumulated_loss, grad_norm, + printf0("step %4d/%d: train loss %f norm %.4f lr %.2e (%.2f ms, %.0f tok/s)\n", + step + 1, train_num_batches, accumulated_loss, grad_norm, step_learning_rate, time_elapsed_ms, bias_corrected_ema_tokens_per_second); logger_log_train(&logger, step, model.mean_loss); @@ -2893,7 +3049,7 @@ int main(int argc, char *argv[]) { // free and destroy everything cudaCheck(cudaEventDestroy(end)); cudaCheck(cudaEventDestroy(start)); - if (hellaswag_available) { evalloader_free(&eval_loader); } + if (run_hellaswag) { evalloader_free(&eval_loader); } dataloader_free(&train_loader); dataloader_free(&val_loader); tokenizer_free(&tokenizer); diff --git a/train_gpt2.py b/train_gpt2.py index f844004..d32ab50 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -47,6 +47,7 @@ class CausalSelfAttention(nn.Module): self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) # output projection self.c_proj = nn.Linear(config.n_embd, config.n_embd) + self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 # regularization self.n_head = config.n_head self.n_embd = config.n_embd @@ -84,6 +85,7 @@ class MLP(nn.Module): self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd) self.gelu = NewGELU() self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd) + self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 def forward(self, x): x = self.c_fc(x) @@ -126,8 +128,27 @@ class GPT(nn.Module): ln_f = nn.LayerNorm(config.n_embd), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) + self.lm_head.LLMC_SKIP_INIT = 1 # don't init this one, we will tie weights self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying + # init all weights, use a torch rng object to be very careful + self.init_rng = torch.Generator() + self.init_rng.manual_seed(42) + self.apply(self._init_weights) + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + # apply special scaled init to the residual projections, per GPT-2 paper + std = 0.02 if not hasattr(module, 'LLMC_RESIDUAL_SCALE_FLAG') else 0.02/math.sqrt(2 * self.config.n_layer) + # we want to skip initializing lm_head, which shares parameters with wte + # and wte was already initialized down below during the Embedding init + if not hasattr(module, 'LLMC_SKIP_INIT'): + torch.nn.init.normal_(module.weight, mean=0.0, std=std, generator=self.init_rng) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02, generator=self.init_rng) + def forward(self, idx, targets=None, return_logits=True): device = idx.device b, t = idx.size() @@ -395,7 +416,7 @@ if __name__ == "__main__": # python train_gpt2.py --inference_only 1 --write_tensors 0 --sequence_length 1024 parser = argparse.ArgumentParser() parser.add_argument("--input_bin", type=str, default="dev/data/tinyshakespeare/tiny_shakespeare_val.bin", help="input .bin to train on") - parser.add_argument("--model", type=str, default="gpt2", help="gpt2|gpt2-medium|gpt2-large|gpt2-xl") + parser.add_argument("--model", type=str, default="gpt2", help="gpt2|gpt2-medium|gpt2-large|gpt2-xl|d12|d24|d36|d48") parser.add_argument("--write_tensors", type=int, default=1, help="write tensors to disk") parser.add_argument("--inference_only", type=int, default=0, help="only run inference") parser.add_argument("--dtype", type=str, default="float32", help="float32|float16|bfloat16") @@ -408,12 +429,12 @@ if __name__ == "__main__": 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_clip", type=float, default=1.0, help="maximum gradient magnitude") + parser.add_argument("--overfit_single_batch", type=int, default=1, help="overfit just one batch of data") args = parser.parse_args() B, T = args.batch_size, args.sequence_length assert 1 <= T <= 1024 assert args.dtype in {"float32", "float16", "bfloat16"} - assert args.model in {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl"} - model_to_size = {"gpt2": "124M", "gpt2-medium": "355M", "gpt2-large": "774M", "gpt2-xl": "1558M"} + assert args.model in {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl", "d12", "d24", "d36", "d48"} # set up DDP (distributed data parallel). torchrun sets this env variable ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run? @@ -427,8 +448,10 @@ if __name__ == "__main__": device = f'cuda:{ddp_local_rank}' torch.cuda.set_device(device) master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. - seed_offset = ddp_rank # each process gets a different seed + seed_offset = 0 # each process gets the exact same seed else: + ddp_rank = 0 + ddp_local_rank = 0 ddp_world_size = 1 master_process = True seed_offset = 0 @@ -449,8 +472,9 @@ if __name__ == "__main__": tokens_per_fwdbwd = B * T * ddp_world_size assert args.total_batch_size % tokens_per_fwdbwd == 0 grad_accum_steps = args.total_batch_size // tokens_per_fwdbwd - print(f"total desired batch size: {args.total_batch_size}") - print(f"=> calculated gradient accumulation steps: {grad_accum_steps}") + if master_process: + print(f"total desired batch size: {args.total_batch_size}") + print(f"=> calculated gradient accumulation steps: {grad_accum_steps}") # set up a context manager following the desired dtype and device ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype] @@ -480,8 +504,19 @@ if __name__ == "__main__": if master_process and args.write_tensors: # tokenizer is technically not tensors but ok write_tokenizer(enc, "gpt2_tokenizer.bin") - # load the GPT-2 model weights - model = GPT.from_pretrained(args.model) + # init the model, either from scratch or from OpenAI pretrained checkpoint + if args.model[0] == "d": + # from scratch (random weights) + model_config = { + "d12": GPTConfig(block_size=1024, vocab_size=50257, n_layer=12, n_head=12, n_embd=768), + "d24": GPTConfig(block_size=1024, vocab_size=50257, n_layer=24, n_head=16, n_embd=1024), + "d36": GPTConfig(block_size=1024, vocab_size=50257, n_layer=36, n_head=20, n_embd=1280), + "d48": GPTConfig(block_size=1024, vocab_size=50257, n_layer=48, n_head=25, n_embd=1600), + }[args.model] + model = GPT(model_config) + else: + # load the GPT-2 model weights + model = GPT.from_pretrained(args.model) model.train() model.to(device) if args.compile: @@ -526,18 +561,20 @@ if __name__ == "__main__": def get_batch(): assert B*T+1 <= len(tokens), "not enough tokens" # for 338,025 tokens. E.g. with B=8 T=1024, this will yield 41 batches before looping - i = 0 + i = B*T*ddp_rank while True: x = tokens[i:i+B*T].view(B, T) y = tokens[i+1:i+B*T+1].view(B, T) yield x, y - i += B*T + i += B*T*ddp_world_size if i + B*T + 1 >= len(tokens): i = 0 # in prod we'd want to randomize the start point a bit + print("We do not expect to reach here in PyTorch right now") + import sys; sys.exit() - # fetch one batch of data, which we will overfit to + # fetch one batch of data data_iter = iter(get_batch()) - x, y = next(data_iter) # we'll overfit this batch below + x, y = next(data_iter) x = x.to(device) y = y.to(device) @@ -549,7 +586,9 @@ if __name__ == "__main__": logits, loss = model(x, y) loss.backward() # save model params, in both float32 and bfloat16 - model_size_str = model_to_size[args.model] # e.g. "124M" + model_to_size = {"gpt2": "124M", "gpt2-medium": "355M", "gpt2-large": "774M", "gpt2-xl": "1558M"} + model_to_size.update({f"d{d}": f"d{d}" for d in [12, 24, 36, 48]}) + model_size_str = model_to_size[args.model] # e.g. "124M", or "d12" write_model(model, f"gpt2_{model_size_str}.bin", dtype="float32") write_model(model, f"gpt2_{model_size_str}_bf16.bin", dtype="bfloat16") # save x, y, logits, loss, and parameter gradients, for debugging C @@ -566,7 +605,7 @@ if __name__ == "__main__": # init the optimizer adam_use_fused = device == "cuda" # only works on CUDA (?) - optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, fused=adam_use_fused) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95), weight_decay=0.0, fused=adam_use_fused) if device == "cuda": torch.cuda.reset_peak_memory_stats() @@ -587,12 +626,17 @@ if __name__ == "__main__": # instead of a SUM we want MEAN, so we scale the loss here loss = loss / grad_accum_steps lossf += loss.item() # keep track of the mean loss + # advance the dataset for the next batch + if not args.overfit_single_batch: + x, y = next(data_iter) + x = x.to(device) + y = y.to(device) + # backward pass if ddp: # we want only the last micro-step to sync grads in a DDP model # the official way to do this is with model.no_sync(), but that is a # context manager that bloats the code, so we just toggle this variable model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1) - # backward pass if not args.inference_only: loss.backward() norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) @@ -608,7 +652,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}, norm: {norm:.3f}") + print0(f"step {step+1:4d}/{args.num_iterations}: train loss {lossf:.6f} norm {norm:.4f} lr 1.00e-04 ({(t1-t0)*1000:.3f} ms, {tokens_per_second:.0f} tok/s)") if step > 0 and step > args.num_iterations - 20: timings.append(t1-t0)