From f470fbd56dcdcd6f81cd5f5a082dba7c0183712f Mon Sep 17 00:00:00 2001 From: ademeure Date: Sat, 20 Jul 2024 15:45:50 +0000 Subject: [PATCH 1/9] Allow restoring weights from the master weights of a checkpoint (deterministically by also saving RNG state of last update) --- llmc/adamw.cuh | 30 +++++++++++++++++++++++++++++- train_gpt2.cu | 47 ++++++++++++++++++++++++++++++++++------------- 2 files changed, 63 insertions(+), 14 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index f806eaa..75ab61a 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -61,6 +61,26 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg ); } +template +__global__ void params_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters, unsigned int seed, bool check_identical) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_parameters) { return; } // guard + + Tp rounded_param; + float param = master_params_memory[idx]; + stochastic_rounding(param, &rounded_param, seed); + + if (check_identical) { + // check if the rounded parameter is identical to the master parameter (for debugging only) + if (params_memory[idx] != rounded_param) { + printf("Mismatch restoring master weights at index %zu (of %zu): %.32f != %.32f\n", idx, num_parameters, (float)params_memory[idx], (float)rounded_param); + assert(false); + } + } + + params_memory[idx] = rounded_param; +} + template void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters, ptrdiff_t w_stride, ptrdiff_t g_stride, ptrdiff_t s_stride, int num_slices, float learning_rate, float beta1, float beta2, int t, float eps, float weight_decay, @@ -75,4 +95,12 @@ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memo learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, grad_scale, seed); cudaCheck(cudaGetLastError()); -} \ No newline at end of file +} + +template +void params_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters, unsigned int seed, bool check_identical=false) { + int block_size = 512; // must match block size of adamw_update so that RNG also matches + int num_blocks = CEIL_DIV(num_parameters, block_size); + params_from_master_kernel<<>>(params_memory, master_params_memory, num_parameters, seed, check_identical); + cudaCheck(cudaGetLastError()); +} diff --git a/train_gpt2.cu b/train_gpt2.cu index 8f11091..8e495e0 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -311,7 +311,8 @@ typedef struct { float* accumulated_mean_loss; // GPU buffer used to accumulate loss across micro-steps float* cpu_losses; // CPU buffer to copy the losses to, allocated with cudaMallocHost unsigned long long rng_state; // the RNG state for seeding stochastic rounding etc. - int use_master_weights; // keep master weights copy in float for optim update? 0|1 + unsigned long long rng_state_last_update; // RNG before last gpt2_update() to re-round identically from master weights + int use_master_weights; // keep master weights copy in float for optim update? 0|1 bool init_state; // set to true if master weights need to be initialized int gelu_fusion; // fuse gelu via cuBLASLt (0=none, 1=forward, 2=forward+backward) int recompute; // recompute gelu | layernorm forward during model backward? 0|1|2 @@ -438,7 +439,7 @@ void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { fcloseCheck(model_file); } -void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { +void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool resuming=false) { if (PRECISION_MODE == PRECISION_FP16) { // TODO for later perhaps, would require us dynamically converting the @@ -483,9 +484,12 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { gpt2_allocate_weights(model); - // read in all the parameters from file and copy them to device - file_to_device(model->params_memory, model_file, model->num_parameters_bytes, - IO_BUF_SIZE, main_stream); + // read in all the parameters from file and copy them to device (if we need them) + if (resuming && model->use_master_weights) { + printf("Resuming from master weights, ignoring model weights in checkpoint...\n"); + } else { + file_to_device(model->params_memory, model_file, model->num_parameters_bytes, IO_BUF_SIZE, main_stream); + } fcloseCheck(model_file); // only return from this function once we are certain the params are ready on the GPU @@ -1008,7 +1012,7 @@ float gpt2_calculate_grad_norm(GPT2 *model, MultiGpuConfig* multi_gpu_config) { return grad_norm_cpu; } -void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, float grad_scale, int t, MultiGpuConfig* multi_gpu_config) { +void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, float grad_scale, int t, MultiGpuConfig* multi_gpu_config, bool init_from_master_only=false) { // update the model parameters using the AdamW optimizer // keep in mind that optimizer sharding (ZeRO-1) assigns different parameters to different GPUs // so we may not be responsible for the entire parameter tensor @@ -1028,6 +1032,10 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo cudaCheck(cudaMemset(model->m_memory, 0, multi_gpu_config->shard_num_parameters * sizeof(float))); cudaCheck(cudaMemset(model->v_memory, 0, multi_gpu_config->shard_num_parameters * sizeof(float))); } + + // save RNG state at this point so we can round from master weights identically when restoring from a checkpoint + model->rng_state_last_update = model->rng_state; + // AdamW update // handle adamw for all the transformer blocks for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { @@ -1064,12 +1072,19 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo cudaCheck(cudaGetLastError()); } - // ok finally call the kernel - adamw_update(param_ptr, master_ptr, grad_ptr, - m_ptr, v_ptr, - shard.size, tensor.size, tensor.size, shard.size, num_layers, - learning_rate, - beta1, beta2, t, eps, wd, grad_scale, seed, main_stream); + if (init_from_master_only) { + // this is only run when resuming training from a checkpoint with master weights + // it allows us to restart training with a different precision amongst other things + assert(master_ptr != NULL); + params_from_master(param_ptr, master_ptr, shard.size, seed, true); + } else { + // ok finally call the kernel to update the weights with AdamW + adamw_update(param_ptr, master_ptr, grad_ptr, + m_ptr, v_ptr, + shard.size, tensor.size, tensor.size, shard.size, num_layers, + learning_rate, + beta1, beta2, t, eps, wd, grad_scale, seed, main_stream); + } cudaCheck(cudaGetLastError()); if (multi_gpu_config->zero_stage == 1) { @@ -1189,6 +1204,7 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) state_header[10] = step; // step of the optimization // model rng state, start at 20 to leave some padding *((unsigned long long*)&state_header[20]) = model->rng_state; // random number generator state + *((unsigned long long*)&state_header[22]) = model->rng_state_last_update; // last gpt2_update // dataloader state, start at 30 to leave some padding *((size_t*)&state_header[30]) = loader->current_shard_idx; // shard of the dataset *((size_t*)&state_header[32]) = loader->current_sample_idx; // position in shard @@ -1225,6 +1241,7 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename int should_shuffle = state_header[5]; // shuffle state of the dataloader *step = state_header[10]; // step of the optimization model->rng_state = *((unsigned long long*)&state_header[20]); // random number generator state + model->rng_state_last_update = *((unsigned long long*)&state_header[22]); // last gpt2_update size_t current_shard_idx = *((size_t*)&state_header[30]); // shard index size_t current_sample_idx = *((size_t*)&state_header[32]); // position in shard @@ -1244,6 +1261,10 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename if(model->use_master_weights) { assert(model->master_weights != nullptr); file_to_device(model->master_weights, state_file, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); + // restore weights from the master weights using the RNG state before last weight update + model->rng_state = model->rng_state_last_update; + gpt2_update(model, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0, &multi_gpu_config, /* init_from_master_only*/ true); + model->rng_state = *((unsigned long long*)&state_header[20]); // use final RNG state from checkpoint after this } model->init_state = false; // we just got the state from file, no need to do first-touch init @@ -1535,7 +1556,7 @@ int main(int argc, char *argv[]) { gpt2_build_from_checkpoint(&model, filename_buffer); } else if (ends_with_bin(load_filename)) { // otherwise, if this is a .bin file, we assume it's a model, let's init from it - gpt2_build_from_checkpoint(&model, load_filename); + gpt2_build_from_checkpoint(&model, load_filename, true); } else { // if it's not .bin, it could be a "special descriptor". This descriptor is used to // construct GPT-2 / GPT-3 models in a convenient format. See the function for docs. From 5cae10f90776562f3074152431a387afceb8fc4a Mon Sep 17 00:00:00 2001 From: ademeure Date: Sat, 20 Jul 2024 16:04:03 +0000 Subject: [PATCH 2/9] make restoring from master weights actually work --- llmc/adamw.cuh | 21 ++++++++++++++------- train_gpt2.cu | 7 ++++--- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index 75ab61a..409c380 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -62,22 +62,27 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg } template -__global__ void params_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters, unsigned int seed, bool check_identical) { +__global__ void params_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters, + ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed, bool check_identical) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= num_parameters) { return; } // guard + if (idx >= num_parameters) { return; } + + // adjust for layer offset + params_memory += blockIdx.y * w_stride; + master_params_memory += blockIdx.y * s_stride; Tp rounded_param; float param = master_params_memory[idx]; stochastic_rounding(param, &rounded_param, seed); if (check_identical) { - // check if the rounded parameter is identical to the master parameter (for debugging only) + // check if the rounded parameter is identical to the master parameter (debugging only) if (params_memory[idx] != rounded_param) { - printf("Mismatch restoring master weights at index %zu (of %zu): %.32f != %.32f\n", idx, num_parameters, (float)params_memory[idx], (float)rounded_param); + printf("Mismatch restoring master weights at index %llu (of %llu): %.20f != %.20f\n", + idx, num_parameters, (float)params_memory[idx], (float)rounded_param); assert(false); } } - params_memory[idx] = rounded_param; } @@ -98,9 +103,11 @@ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memo } template -void params_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters, unsigned int seed, bool check_identical=false) { +void params_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters, + ptrdiff_t w_stride, ptrdiff_t s_stride, int num_slices, unsigned int seed, cudaStream_t stream, bool check_identical=false) { int block_size = 512; // must match block size of adamw_update so that RNG also matches int num_blocks = CEIL_DIV(num_parameters, block_size); - params_from_master_kernel<<>>(params_memory, master_params_memory, num_parameters, seed, check_identical); + params_from_master_kernel<<>> + (params_memory, master_params_memory, num_parameters, w_stride, s_stride, seed, check_identical); cudaCheck(cudaGetLastError()); } diff --git a/train_gpt2.cu b/train_gpt2.cu index 8e495e0..99f4b0a 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1076,7 +1076,8 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo // this is only run when resuming training from a checkpoint with master weights // it allows us to restart training with a different precision amongst other things assert(master_ptr != NULL); - params_from_master(param_ptr, master_ptr, shard.size, seed, true); + params_from_master(param_ptr, master_ptr, + shard.size, tensor.size, shard.size, num_layers, seed, main_stream); } else { // ok finally call the kernel to update the weights with AdamW adamw_update(param_ptr, master_ptr, grad_ptr, @@ -1553,10 +1554,10 @@ int main(int argc, char *argv[]) { gpt2_init_common(&model); if (resuming == 1) { // if `-y 1` was set, then we are resuming from the latest checkpoint - gpt2_build_from_checkpoint(&model, filename_buffer); + gpt2_build_from_checkpoint(&model, filename_buffer, true); } else if (ends_with_bin(load_filename)) { // otherwise, if this is a .bin file, we assume it's a model, let's init from it - gpt2_build_from_checkpoint(&model, load_filename, true); + gpt2_build_from_checkpoint(&model, load_filename); } else { // if it's not .bin, it could be a "special descriptor". This descriptor is used to // construct GPT-2 / GPT-3 models in a convenient format. See the function for docs. From 9781627a8aae98b346c042d244bbd9d92bfdce4e Mon Sep 17 00:00:00 2001 From: ademeure Date: Sat, 20 Jul 2024 16:06:24 +0000 Subject: [PATCH 3/9] allow restoring from checkpoint of different precision --- train_gpt2.cu | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 99f4b0a..d926ae0 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -462,16 +462,20 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool r fprintf(stderr, "---> HINT: try to re-run `python train_gpt2.py`\n"); exit(EXIT_FAILURE); } - if (PRECISION_MODE == PRECISION_BF16 && version != 5) { - fprintf(stderr, "Precision is configured as BF16 but model at %s is not.\n", checkpoint_path); - fprintf(stderr, "---> HINT: are you sure you're loading a _bf16.bin file?\n"); - exit(EXIT_FAILURE); - } - if (PRECISION_MODE == PRECISION_FP32 && version != 3) { - fprintf(stderr, "Precision is configured as FP32 but model at %s is not.\n", checkpoint_path); - fprintf(stderr, "---> HINT: to turn on FP32 you have to compile like: `make train_gpt2cu PRECISION=FP32`\n"); - fprintf(stderr, "---> HINT: are you sure you're loading a .bin file without any _bf16 in the name?\n"); - exit(EXIT_FAILURE); + + // check if the precision mode matches the model (don't care if restoring from master weights!) + if (!resuming || !model->use_master_weights) { + if (PRECISION_MODE == PRECISION_BF16 && version != 5) { + fprintf(stderr, "Precision is configured as BF16 but model at %s is not.\n", checkpoint_path); + fprintf(stderr, "---> HINT: are you sure you're loading a _bf16.bin file?\n"); + exit(EXIT_FAILURE); + } + if (PRECISION_MODE == PRECISION_FP32 && version != 3) { + fprintf(stderr, "Precision is configured as FP32 but model at %s is not.\n", checkpoint_path); + fprintf(stderr, "---> HINT: to turn on FP32 you have to compile like: `make train_gpt2cu PRECISION=FP32`\n"); + fprintf(stderr, "---> HINT: are you sure you're loading a .bin file without any _bf16 in the name?\n"); + exit(EXIT_FAILURE); + } } // read in hyperparameters From 2eabc223cc549fcd1a556ae3efd4b4cf54754ce6 Mon Sep 17 00:00:00 2001 From: ademeure Date: Sat, 20 Jul 2024 16:31:27 +0000 Subject: [PATCH 4/9] simplify a little bit --- llmc/adamw.cuh | 14 ++------------ train_gpt2.cu | 8 +++----- 2 files changed, 5 insertions(+), 17 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index 409c380..6d64438 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -66,23 +66,13 @@ __global__ void params_from_master_kernel(Tp* params_memory, float* master_param ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed, bool check_identical) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_parameters) { return; } - - // adjust for layer offset - params_memory += blockIdx.y * w_stride; + params_memory += blockIdx.y * w_stride; // adjust for layer offset master_params_memory += blockIdx.y * s_stride; Tp rounded_param; float param = master_params_memory[idx]; stochastic_rounding(param, &rounded_param, seed); - - if (check_identical) { - // check if the rounded parameter is identical to the master parameter (debugging only) - if (params_memory[idx] != rounded_param) { - printf("Mismatch restoring master weights at index %llu (of %llu): %.20f != %.20f\n", - idx, num_parameters, (float)params_memory[idx], (float)rounded_param); - assert(false); - } - } + assert(!check_identical || params_memory[idx] == rounded_param); // for debugging only (needs non-master params to be loaded as well) params_memory[idx] = rounded_param; } diff --git a/train_gpt2.cu b/train_gpt2.cu index d926ae0..4ae2de8 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1077,11 +1077,9 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo } if (init_from_master_only) { - // this is only run when resuming training from a checkpoint with master weights - // it allows us to restart training with a different precision amongst other things - assert(master_ptr != NULL); - params_from_master(param_ptr, master_ptr, - shard.size, tensor.size, shard.size, num_layers, seed, main_stream); + // when resuming training from a checkpoint with master weights + init_from_master(param_ptr, master_ptr, + shard.size, tensor.size, shard.size, num_layers, seed, main_stream); } else { // ok finally call the kernel to update the weights with AdamW adamw_update(param_ptr, master_ptr, grad_ptr, From 4d77ece6b4f15fd41e2c927ad4838e5bd750e5bb Mon Sep 17 00:00:00 2001 From: ademeure Date: Sat, 20 Jul 2024 16:45:24 +0000 Subject: [PATCH 5/9] simplified further (don't need non-functional error checking...) --- llmc/adamw.cuh | 13 ++++--------- train_gpt2.cu | 9 ++++----- 2 files changed, 8 insertions(+), 14 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index 6d64438..e15e81a 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -62,18 +62,13 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg } template -__global__ void params_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters, +__global__ void init_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters, ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed, bool check_identical) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_parameters) { return; } params_memory += blockIdx.y * w_stride; // adjust for layer offset master_params_memory += blockIdx.y * s_stride; - - Tp rounded_param; - float param = master_params_memory[idx]; - stochastic_rounding(param, &rounded_param, seed); - assert(!check_identical || params_memory[idx] == rounded_param); // for debugging only (needs non-master params to be loaded as well) - params_memory[idx] = rounded_param; + stochastic_rounding(master_params_memory[idx], ¶ms_memory[idx], seed); } template @@ -93,11 +88,11 @@ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memo } template -void params_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters, +void init_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters, ptrdiff_t w_stride, ptrdiff_t s_stride, int num_slices, unsigned int seed, cudaStream_t stream, bool check_identical=false) { int block_size = 512; // must match block size of adamw_update so that RNG also matches int num_blocks = CEIL_DIV(num_parameters, block_size); - params_from_master_kernel<<>> + init_from_master_kernel<<>> (params_memory, master_params_memory, num_parameters, w_stride, s_stride, seed, check_identical); cudaCheck(cudaGetLastError()); } diff --git a/train_gpt2.cu b/train_gpt2.cu index 4ae2de8..61a63b6 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1016,7 +1016,8 @@ float gpt2_calculate_grad_norm(GPT2 *model, MultiGpuConfig* multi_gpu_config) { return grad_norm_cpu; } -void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, float grad_scale, int t, MultiGpuConfig* multi_gpu_config, bool init_from_master_only=false) { +void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, float grad_scale, int t, + MultiGpuConfig* multi_gpu_config, bool init_from_master_only=false) { // update the model parameters using the AdamW optimizer // keep in mind that optimizer sharding (ZeRO-1) assigns different parameters to different GPUs // so we may not be responsible for the entire parameter tensor @@ -1077,9 +1078,8 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo } if (init_from_master_only) { - // when resuming training from a checkpoint with master weights - init_from_master(param_ptr, master_ptr, - shard.size, tensor.size, shard.size, num_layers, seed, main_stream); + // when resuming training from a checkpoint with master weights (allows changing precision) + init_from_master(param_ptr, master_ptr, shard.size, tensor.size, shard.size, num_layers, seed, main_stream); } else { // ok finally call the kernel to update the weights with AdamW adamw_update(param_ptr, master_ptr, grad_ptr, @@ -1088,7 +1088,6 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo learning_rate, beta1, beta2, t, eps, wd, grad_scale, seed, main_stream); } - cudaCheck(cudaGetLastError()); if (multi_gpu_config->zero_stage == 1) { #if MULTI_GPU From 52e6e0f80b4dad84a772aaf86deee759738e3361 Mon Sep 17 00:00:00 2001 From: ademeure Date: Thu, 25 Jul 2024 21:32:01 +0000 Subject: [PATCH 6/9] fix bug from merge (init_state set to false too late) --- train_gpt2.cu | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 61a63b6..c1696eb 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -489,9 +489,8 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool r gpt2_allocate_weights(model); // read in all the parameters from file and copy them to device (if we need them) - if (resuming && model->use_master_weights) { - printf("Resuming from master weights, ignoring model weights in checkpoint...\n"); - } else { + // if we are restoring with master weights, ignore these weights and init from master instead + if (!resuming || !model->use_master_weights) { file_to_device(model->params_memory, model_file, model->num_parameters_bytes, IO_BUF_SIZE, main_stream); } fcloseCheck(model_file); @@ -1256,6 +1255,8 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename printf0("Error: Master weights requested, but not present in state file."); exit(EXIT_FAILURE); } + + model->init_state = false; // we just got the state from file, no need to do first-touch init assert(model->m_memory != nullptr); assert(model->v_memory != nullptr); file_to_device(model->m_memory, state_file, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); @@ -1269,8 +1270,6 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename model->rng_state = *((unsigned long long*)&state_header[20]); // use final RNG state from checkpoint after this } - model->init_state = false; // we just got the state from file, no need to do first-touch init - // revive the DataLoader object and its state loader->should_shuffle = should_shuffle; if (should_shuffle == 1) { From a794bcb395c459b1173d1b2bcf99321b6bbd6851 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 28 Jul 2024 17:10:20 +0000 Subject: [PATCH 7/9] small fixes --- llmc/adamw.cuh | 6 +++--- train_gpt2.cu | 33 +++++++++++++++++++-------------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index e15e81a..4453576 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -63,7 +63,7 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg template __global__ void init_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters, - ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed, bool check_identical) { + ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_parameters) { return; } params_memory += blockIdx.y * w_stride; // adjust for layer offset @@ -89,10 +89,10 @@ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memo template void init_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters, - ptrdiff_t w_stride, ptrdiff_t s_stride, int num_slices, unsigned int seed, cudaStream_t stream, bool check_identical=false) { + ptrdiff_t w_stride, ptrdiff_t s_stride, int num_slices, unsigned int seed, cudaStream_t stream) { int block_size = 512; // must match block size of adamw_update so that RNG also matches int num_blocks = CEIL_DIV(num_parameters, block_size); init_from_master_kernel<<>> - (params_memory, master_params_memory, num_parameters, w_stride, s_stride, seed, check_identical); + (params_memory, master_params_memory, num_parameters, w_stride, s_stride, seed); cudaCheck(cudaGetLastError()); } diff --git a/train_gpt2.cu b/train_gpt2.cu index 5d1dce5..bea3d57 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -439,7 +439,11 @@ void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { fcloseCheck(model_file); } -void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool resuming=false) { +void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool weight_init=true) { + // If weight_init is true, we will load the weights from this checkpoint .bin file + // We sometimes want this to be false, if we are going to initialize these weights from + // the master weights that are instead stored in the state .bin file. + // In that case, this function mostly loads the model hyperparameters from the header. if (PRECISION_MODE == PRECISION_FP16) { // TODO for later perhaps, would require us dynamically converting the @@ -463,8 +467,8 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool r exit(EXIT_FAILURE); } - // check if the precision mode matches the model (don't care if restoring from master weights!) - if (!resuming || !model->use_master_weights) { + // check if the precision mode of the checkpoing matches the model precision + if (weight_init) { if (PRECISION_MODE == PRECISION_BF16 && version != 5) { fprintf(stderr, "Precision is configured as BF16 but model at %s is not.\n", checkpoint_path); fprintf(stderr, "---> HINT: are you sure you're loading a _bf16.bin file?\n"); @@ -486,11 +490,10 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool r model->config.channels = model_header[6]; model->config.padded_vocab_size = model_header[7]; - gpt2_allocate_weights(model); - - // read in all the parameters from file and copy them to device (if we need them) - // if we are restoring with master weights, ignore these weights and init from master instead - if (!resuming || !model->use_master_weights) { + // read in the parameters if weight_init is true + // we are assuming the space has already been allocated! + if (weight_init) { + assert(model->params_memory != NULL); file_to_device(model->params_memory, model_file, model->num_parameters_bytes, IO_BUF_SIZE, main_stream); } fcloseCheck(model_file); @@ -1538,13 +1541,12 @@ int main(int argc, char *argv[]) { // figure out if we are going to be resuming the optimization int resuming = 0; + // find the DONE file with the highest step count int resume_max_step = find_max_step(output_log_dir); - if (resume == 1) { - // find the DONE file with the highest step count + if (resume == 1) { // is -y 1 resume flag set? assert(output_log_dir != NULL); - if (resume_max_step == -1) { - } else { - resuming = 1; + if (resume_max_step != -1) { + resuming = 1; // -y 1 is set, and we found a checkpoint we can resume from snprintf(filename_buffer, sizeof(filename_buffer), "%s/model_%08d.bin", output_log_dir, resume_max_step); } } @@ -1552,9 +1554,12 @@ int main(int argc, char *argv[]) { // build the GPT-2 model GPT2 model; gpt2_init_common(&model); + gpt2_allocate_weights(&model); if (resuming == 1) { // if `-y 1` was set, then we are resuming from the latest checkpoint - gpt2_build_from_checkpoint(&model, filename_buffer, true); + // if we are using master weights, we'll init them later inside load_state() + bool weight_init = !use_master_weights; + gpt2_build_from_checkpoint(&model, filename_buffer, weight_init); } else if (ends_with_bin(load_filename)) { // otherwise, if this is a .bin file, we assume it's a model, let's init from it gpt2_build_from_checkpoint(&model, load_filename); From 2b827f1659f1f6146d600dbbf512ae9bbe7ebeab Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Mon, 29 Jul 2024 02:39:50 +0000 Subject: [PATCH 8/9] bring back state allocation into build_from_checkpoint --- train_gpt2.cu | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index bea3d57..2acd00d 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -490,6 +490,9 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool w model->config.channels = model_header[6]; model->config.padded_vocab_size = model_header[7]; + // allocate memory for the model parameters + gpt2_allocate_weights(model); + // read in the parameters if weight_init is true // we are assuming the space has already been allocated! if (weight_init) { @@ -1554,7 +1557,6 @@ int main(int argc, char *argv[]) { // build the GPT-2 model GPT2 model; gpt2_init_common(&model); - gpt2_allocate_weights(&model); if (resuming == 1) { // if `-y 1` was set, then we are resuming from the latest checkpoint // if we are using master weights, we'll init them later inside load_state() From 51dd102328a68d3a53b6588b09cf3f60dfd1bb36 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Tue, 30 Jul 2024 20:01:43 +0000 Subject: [PATCH 9/9] remove confusing comment --- train_gpt2.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 2acd00d..2737e64 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -494,7 +494,6 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path, bool w gpt2_allocate_weights(model); // read in the parameters if weight_init is true - // we are assuming the space has already been allocated! if (weight_init) { assert(model->params_memory != NULL); file_to_device(model->params_memory, model_file, model->num_parameters_bytes, IO_BUF_SIZE, main_stream);