From 3c3c965840b239480c87d3dfe04a70bf70986164 Mon Sep 17 00:00:00 2001 From: Azret Botash Date: Tue, 14 May 2024 17:16:24 -0700 Subject: [PATCH 01/14] Adding Mersenne Twisters C --- rand.h | 141 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 rand.h diff --git a/rand.h b/rand.h new file mode 100644 index 0000000..f69340d --- /dev/null +++ b/rand.h @@ -0,0 +1,141 @@ +#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 From d09631807ac0c594891d163a9129c381448cab13 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 14:35:46 +0000 Subject: [PATCH 02/14] first draft of random init, crashes with some cuBLAS error, debugging --- train_gpt2.cu | 96 +++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 94 insertions(+), 2 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 798fe76..31409a1 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2035,6 +2035,84 @@ 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 = 12; num_heads = 12; } // gpt2 (124M) + else if (depth == 24) { channels = 16; num_heads = 16; } // gpt2-medium (350M) + else if (depth == 36) { channels = 20; num_heads = 20; } // gpt2-large (774M) + else if (depth == 48) { channels = 25; 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 + unsigned long long init_rng_state = 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 + size_t offset = 0; + float residual_scale = 1.0f / sqrtf(2.0f * model->config.num_layers); + for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + // hard-coding the positions of the weights tensors here + if (i == 0 || i == 1 || i == 2 || i == 4 || i == 6 || i == 8 || i == 10 + || i == 12 || i == 14) { + // 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; + for (size_t j = 0; j < model->param_elements[i]; j++) { + float f = random_f32(&init_rng_state); // random float in [0, 1] + f *= scale; + f -= 0.5f * scale; // mean 0 + params_memory_cpu[offset + j] = (floatX)f; + } + } + 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 @@ -2666,9 +2744,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); From e6a7d1d3e9de3b65fbb43d21b4e4590a9718a2c4 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 14:59:25 +0000 Subject: [PATCH 03/14] allow the python script to also init from random and save those weights, so it's a good reference for our C implementation --- train_gpt2.py | 39 +++++++++++++++++++++++++++++++++------ 1 file changed, 33 insertions(+), 6 deletions(-) diff --git a/train_gpt2.py b/train_gpt2.py index f844004..ce1c254 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -128,6 +128,21 @@ class GPT(nn.Module): self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying + # init all weights + self.apply(self._init_weights) + # apply special scaled init to the residual projections, per GPT-2 paper + for pn, p in self.named_parameters(): + if pn.endswith('c_proj.weight'): + torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer)) + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + 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) + def forward(self, idx, targets=None, return_logits=True): device = idx.device b, t = idx.size() @@ -395,7 +410,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") @@ -412,8 +427,7 @@ if __name__ == "__main__": 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? @@ -480,8 +494,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: @@ -549,7 +574,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 From 70a9c75348dee80def7172e96326522abcbdf86b Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 15:16:49 +0000 Subject: [PATCH 04/14] use pytorch rand and fix dumb bug lol --- rand.h | 84 ++++++++++++++++++++++++++++++++++++++++++++++++++- train_gpt2.cu | 26 +++++++++------- 2 files changed, 98 insertions(+), 12 deletions(-) diff --git a/rand.h b/rand.h index f69340d..e60e5e6 100644 --- a/rand.h +++ b/rand.h @@ -1,3 +1,85 @@ +/* +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 @@ -81,7 +163,7 @@ void uniform_(float* data, unsigned int numel, float from, float to, mt19937_sta } } -// Box–Muller transform +// Box�Muller transform void normal_fill_16(float* data, float mean, float std, mt19937_state* state) { #define EPSILONE 1e-12 diff --git a/train_gpt2.cu b/train_gpt2.cu index 31409a1..edd3165 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 @@ -2042,10 +2044,10 @@ void gpt2_build_from_random(GPT2 *model, int depth) { model->config.num_layers = depth; // follows GPT-2 sizes int channels, num_heads; - if (depth == 12) { channels = 12; num_heads = 12; } // gpt2 (124M) - else if (depth == 24) { channels = 16; num_heads = 16; } // gpt2-medium (350M) - else if (depth == 36) { channels = 20; num_heads = 20; } // gpt2-large (774M) - else if (depth == 48) { channels = 25; num_heads = 25; } // gpt2-xl (1558M) + 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; @@ -2068,7 +2070,8 @@ void gpt2_build_from_random(GPT2 *model, int depth) { // 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 - unsigned long long init_rng_state = 42; + mt19937_state rng_state; + manual_seed(&rng_state, 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 @@ -2081,12 +2084,13 @@ void gpt2_build_from_random(GPT2 *model, int depth) { // 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; - for (size_t j = 0; j < model->param_elements[i]; j++) { - float f = random_f32(&init_rng_state); // random float in [0, 1] - f *= scale; - f -= 0.5f * scale; // mean 0 - params_memory_cpu[offset + j] = (floatX)f; + int n = model->param_elements[i]; + float *fp32_buffer = (float*)mallocCheck(n * sizeof(float)); + normal_(fp32_buffer, n, 0.0f, scale, &rng_state); + for (size_t j = 0; j < n; j++) { + params_memory_cpu[offset + j] = (floatX)fp32_buffer[j]; } + free(fp32_buffer); } offset += model->param_elements[i]; } From 86682af9a9286714eb7879047ca59100a4cd485f Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 17:50:33 +0000 Subject: [PATCH 05/14] llm.c matches pytorch init from scratch exactly now --- train_gpt2.cu | 60 +++++++++++++++++++++++++++++++++++---------------- train_gpt2.py | 20 +++++++++++------ 2 files changed, 55 insertions(+), 25 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index edd3165..8fd0093 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2063,37 +2063,61 @@ void gpt2_build_from_random(GPT2 *model, int depth) { 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 rng_state; - manual_seed(&rng_state, 42); + 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 - size_t offset = 0; float residual_scale = 1.0f / sqrtf(2.0f * model->config.num_layers); - for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { - // hard-coding the positions of the weights tensors here - if (i == 0 || i == 1 || i == 2 || i == 4 || i == 6 || i == 8 || i == 10 - || i == 12 || i == 14) { - // 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; - int n = model->param_elements[i]; - float *fp32_buffer = (float*)mallocCheck(n * sizeof(float)); - normal_(fp32_buffer, n, 0.0f, scale, &rng_state); - for (size_t j = 0; j < n; j++) { - params_memory_cpu[offset + j] = (floatX)fp32_buffer[j]; + // 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; + } } - free(fp32_buffer); + // 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]; } - 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); diff --git a/train_gpt2.py b/train_gpt2.py index ce1c254..66a9bbd 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,22 +128,26 @@ 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 + # 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) - # apply special scaled init to the residual projections, per GPT-2 paper - for pn, p in self.named_parameters(): - if pn.endswith('c_proj.weight'): - torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer)) def _init_weights(self, module): if isinstance(module, nn.Linear): - torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + # 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) + 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 From 1f91bfc44206dfe89bef67ab353f81ed75f40568 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 18:19:25 +0000 Subject: [PATCH 06/14] fix small bug on eval logging --- train_gpt2.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 8fd0093..585e82a 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2635,9 +2635,9 @@ void logger_init(Logger *logger, const char *filename) { if (filename != NULL) { 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); } } @@ -2900,7 +2900,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 From 5f87b13f34de9bc5e2698d1bec2347f88ceb67ba Mon Sep 17 00:00:00 2001 From: otabuzzman Date: Thu, 23 May 2024 21:15:01 +0200 Subject: [PATCH 07/14] Update documentation with Swift port reference --- README.md | 3 +++ 1 file changed, 3 insertions(+) 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 From 3cb2812774b344af1224197580af55fee2063683 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 19:27:04 +0000 Subject: [PATCH 08/14] skip hellaswag eval on step 0 i think... not sure but ok for now --- train_gpt2.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 585e82a..1a025f2 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2886,7 +2886,7 @@ int main(int argc, char *argv[]) { // once in a while estimate HellaSwag accuracy if (hellaswag_available && - (step % val_loss_every == 0 || last_step)) { + ((step > 0 && step % val_loss_every == 0) || last_step)) { NvtxRange evaluation_range("evaluation"); float eval_acc_norm = 0.0f; evalloader_reset(&eval_loader); From 949d71a3d22869e55d87ac2d927fcf2dcabc09f3 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 20:32:56 +0000 Subject: [PATCH 09/14] only rank 0 logs --- train_gpt2.cu | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 1a025f2..5903815 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2630,9 +2630,12 @@ 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) { @@ -2650,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); } } } @@ -2844,9 +2847,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"); From 645869b6f7f670c84e4a5d6c4945fa0f8eeb1338 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 20:45:59 +0000 Subject: [PATCH 10/14] add weight decay -c option and be more careful in our tests of correctness, our weight decay didn't match to pytorch. also modify the betas in AdamW to be consistent with those used in GPT-3 training --- test_gpt2.cu | 18 +++++++++--------- train_gpt2.cu | 6 +++++- train_gpt2.py | 2 +- 3 files changed, 15 insertions(+), 11 deletions(-) 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 5903815..65afa74 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2675,6 +2675,7 @@ 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, " -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"); @@ -2702,6 +2703,7 @@ 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; + 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? @@ -2726,6 +2728,7 @@ 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] == '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]); } @@ -2754,6 +2757,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("| 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); @@ -2982,7 +2986,7 @@ 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); + float grad_norm = gpt2_update(&model, 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); diff --git a/train_gpt2.py b/train_gpt2.py index 66a9bbd..298ab70 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -599,7 +599,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() From 661975cc3c6a0e1684dcc91430f4e0fb4ebbdf23 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 20:57:04 +0000 Subject: [PATCH 11/14] add learning rate warmup option --- train_gpt2.cu | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 65afa74..5071390 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2675,6 +2675,7 @@ 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, " -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"); @@ -2703,6 +2704,7 @@ 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 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? @@ -2728,6 +2730,7 @@ 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] == '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]); } @@ -2757,6 +2760,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("| warmup iterations | %-50d |\n", warmup_iterations); printf0("| weight decay | %-50e |\n", weight_decay); printf0("| grad_clip | %-50e |\n", grad_clip); printf0("| max_steps | %-50d |\n", max_steps); @@ -2986,7 +2990,14 @@ 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.95f, 1e-8f, weight_decay, grad_clip, step+1, &multi_gpu_config); + // learning rate schedule + float step_learning_rate = learning_rate; + if (warmup_iterations > 0) { + float lr_scale = fminf(1.0f, (float)(step + 1) / warmup_iterations); + step_learning_rate *= lr_scale; + } + // 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); @@ -3008,8 +3019,8 @@ int main(int argc, char *argv[]) { 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); From 64b6a146798d7e03c2354ecb599eaa8c91820551 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 23 May 2024 21:19:00 +0000 Subject: [PATCH 12/14] add learning rate decay schedule, now we have the full scheduler implemented --- train_gpt2.cu | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 5071390..ce1fa30 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2676,6 +2676,7 @@ void error_usage() { 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"); @@ -2705,6 +2706,7 @@ int main(int argc, char *argv[]) { 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? @@ -2731,6 +2733,7 @@ int main(int argc, char *argv[]) { 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]); } @@ -2745,6 +2748,8 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 'r') { recompute = 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 @@ -2759,8 +2764,9 @@ 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); @@ -2990,11 +2996,17 @@ int main(int argc, char *argv[]) { model.mean_loss = lossf; // update the parameters gpt2_multi_gpu_accumulate(&model, &multi_gpu_config); - // learning rate schedule + // learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac float step_learning_rate = learning_rate; - if (warmup_iterations > 0) { - float lr_scale = fminf(1.0f, (float)(step + 1) / warmup_iterations); - step_learning_rate *= lr_scale; + 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); From 032e76c25974855b3c9282e07b98c232082d314b Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 24 May 2024 01:02:21 +0000 Subject: [PATCH 13/14] start putting llm.c and pytorch right next to each other, identical training runs with identical results and prints. almost --- train_gpt2.cu | 4 ++-- train_gpt2.py | 29 ++++++++++++++++++++--------- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index ce1fa30..54a0c4f 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -3022,10 +3022,10 @@ 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)); diff --git a/train_gpt2.py b/train_gpt2.py index 298ab70..d32ab50 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -429,6 +429,7 @@ 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 @@ -447,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 @@ -469,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] @@ -557,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) @@ -620,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) @@ -641,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) From dee4e42548fa7e14ddd16f45479067300173fd0e Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 24 May 2024 03:47:07 +0000 Subject: [PATCH 14/14] add option to not run hellaswag, interferes with a bunch of testing, e.g. if T is low --- train_gpt2.cu | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 54a0c4f..81f82f7 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2688,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); } @@ -2719,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 @@ -2746,6 +2748,7 @@ 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 @@ -2832,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 @@ -2847,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); @@ -2904,7 +2908,7 @@ int main(int argc, char *argv[]) { } // once in a while estimate HellaSwag accuracy - if (hellaswag_available && + if (run_hellaswag && ((step > 0 && step % val_loss_every == 0) || last_step)) { NvtxRange evaluation_range("evaluation"); float eval_acc_norm = 0.0f; @@ -3045,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);