From 04234d0c24e97b68d69844bb72d9f895c81e50e4 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Thu, 6 Jun 2024 00:49:11 +0300 Subject: [PATCH 01/47] utility functions for device <-> disk IO --- llmc/cuda_common.h | 81 ++++++++++++++++++++++++++++++++++++++++++ llmc/utils.h | 36 +++++++++++++++---- test/device_file_io.cu | 53 +++++++++++++++++++++++++++ 3 files changed, 163 insertions(+), 7 deletions(-) create mode 100644 test/device_file_io.cu diff --git a/llmc/cuda_common.h b/llmc/cuda_common.h index 3918bf5..921bf23 100644 --- a/llmc/cuda_common.h +++ b/llmc/cuda_common.h @@ -15,6 +15,8 @@ Common utilities for CUDA code. #include #include +#include "utils.h" + // ---------------------------------------------------------------------------- // Global defines and settings @@ -116,4 +118,83 @@ class NvtxRange { }; #define NVTX_RANGE_FN() NvtxRange nvtx_range(__FUNCTION__) +// copy num_bytes from device pointer src into file dest, using double buffering running on the given stream. +inline void device_to_file(FILE* dest, void* src, size_t num_bytes, size_t buffer_size, cudaStream_t stream) { + // allocate pinned buffer for faster, async transfer + char* buffer_space; + cudaCheck(cudaMallocHost(&buffer_space, 2*buffer_size)); + // split allocation in two + void* read_buffer = buffer_space; + void* write_buffer = buffer_space + buffer_size; + + // prime the read buffer; first copy means we have to wait + char* gpu_read_ptr = (char*)src; + size_t copy_amount = std::min(buffer_size, num_bytes); + cudaCheck(cudaMemcpyAsync(read_buffer, gpu_read_ptr, copy_amount, cudaMemcpyDeviceToHost, stream)); + cudaCheck(cudaStreamSynchronize(stream)); + size_t rest_bytes = num_bytes - copy_amount; + size_t write_buffer_size = copy_amount; + gpu_read_ptr += copy_amount; + + std::swap(read_buffer, write_buffer); + // now the main loop; as long as there are bytes left + while(rest_bytes > 0) { + // initiate next read + copy_amount = std::min(buffer_size, rest_bytes); + cudaCheck(cudaMemcpyAsync(read_buffer, gpu_read_ptr, copy_amount, cudaMemcpyDeviceToHost, stream)); + // while this is going on, transfer the write buffer to disk + fwriteCheck(write_buffer, 1, write_buffer_size, dest); + cudaCheck(cudaStreamSynchronize(stream)); // wait for both buffers to be ready. + + std::swap(read_buffer, write_buffer); + rest_bytes -= copy_amount; + write_buffer_size = copy_amount; + gpu_read_ptr += copy_amount; + } + + // make sure to write the last remaining write buffer + fwriteCheck(write_buffer, 1, write_buffer_size, dest); + cudaCheck(cudaFreeHost(buffer_space)); +} + +// copy num_bytes from file src into device pointer dest, using double buffering running on the given stream. +inline void file_to_device(void* dest, FILE* src, size_t num_bytes, size_t buffer_size, cudaStream_t stream) { + // allocate pinned buffer for faster, async transfer + // from the docs (https://developer.download.nvidia.com/compute/DevZone/docs/html/C/doc/html/group__CUDART__HIGHLEVEL_ge439496de696b166ba457dab5dd4f356.html) + // WC memory is a good option for buffers that will be written by the CPU and read by the device via mapped pinned memory or host->device transfers. + char* buffer_space; + cudaCheck(cudaMallocHost(&buffer_space, 2*buffer_size, cudaHostAllocWriteCombined)); + // split allocation in two + void* read_buffer = buffer_space; + void* write_buffer = buffer_space + buffer_size; + + // prime the read buffer; + char* gpu_write_ptr = (char*)dest; + size_t copy_amount = std::min(buffer_size, num_bytes); + freadCheck(read_buffer, 1, copy_amount, src); + + size_t rest_bytes = num_bytes - copy_amount; + size_t write_buffer_size = copy_amount; + std::swap(read_buffer, write_buffer); + + // now the main loop; as long as there are bytes left + while(rest_bytes > 0) { + // initiate next read + copy_amount = std::min(buffer_size, rest_bytes); + cudaCheck(cudaMemcpyAsync(gpu_write_ptr, write_buffer, write_buffer_size, cudaMemcpyHostToDevice, stream)); + gpu_write_ptr += write_buffer_size; + // while this is going on, read from disk + freadCheck(read_buffer, 1, copy_amount, src); + cudaCheck(cudaStreamSynchronize(stream)); // wait for both buffers to be ready. + + std::swap(read_buffer, write_buffer); + rest_bytes -= copy_amount; + write_buffer_size = copy_amount; + } + + // copy the last remaining write buffer to gpu + cudaCheck(cudaMemcpyAsync(gpu_write_ptr, write_buffer, write_buffer_size, cudaMemcpyHostToDevice, stream)); + cudaCheck(cudaFreeHost(buffer_space)); +} + #endif // CUDA_COMMON_H \ No newline at end of file diff --git a/llmc/utils.h b/llmc/utils.h index be8acdb..e533f2b 100644 --- a/llmc/utils.h +++ b/llmc/utils.h @@ -21,7 +21,7 @@ // simple replace fopen, fread, fclose, fseek // with fopenCheck, freadCheck, fcloseCheck, fseekCheck -FILE *fopen_check(const char *path, const char *mode, const char *file, int line) { +inline FILE *fopen_check(const char *path, const char *mode, const char *file, int line) { FILE *fp = fopen(path, mode); if (fp == NULL) { fprintf(stderr, "Error: Failed to open file '%s' at %s:%d\n", path, file, line); @@ -39,7 +39,7 @@ FILE *fopen_check(const char *path, const char *mode, const char *file, int line #define fopenCheck(path, mode) fopen_check(path, mode, __FILE__, __LINE__) -void fread_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { +inline void fread_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { size_t result = fread(ptr, size, nmemb, stream); if (result != nmemb) { if (feof(stream)) { @@ -61,7 +61,7 @@ void fread_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char #define freadCheck(ptr, size, nmemb, stream) fread_check(ptr, size, nmemb, stream, __FILE__, __LINE__) -void fclose_check(FILE *fp, const char *file, int line) { +inline void fclose_check(FILE *fp, const char *file, int line) { if (fclose(fp) != 0) { fprintf(stderr, "Error: Failed to close file at %s:%d\n", file, line); fprintf(stderr, "Error details:\n"); @@ -73,7 +73,7 @@ void fclose_check(FILE *fp, const char *file, int line) { #define fcloseCheck(fp) fclose_check(fp, __FILE__, __LINE__) -void fseek_check(FILE *fp, long off, int whence, const char *file, int line) { +inline void fseek_check(FILE *fp, long off, int whence, const char *file, int line) { if (fseek(fp, off, whence) != 0) { fprintf(stderr, "Error: Failed to seek in file at %s:%d\n", file, line); fprintf(stderr, "Error details:\n"); @@ -87,10 +87,32 @@ void fseek_check(FILE *fp, long off, int whence, const char *file, int line) { #define fseekCheck(fp, off, whence) fseek_check(fp, off, whence, __FILE__, __LINE__) +inline void fwrite_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { + size_t result = fwrite(ptr, size, nmemb, stream); + if (result != nmemb) { + if (feof(stream)) { + fprintf(stderr, "Error: Unexpected end of file at %s:%d\n", file, line); + } else if (ferror(stream)) { + fprintf(stderr, "Error: File write error at %s:%d\n", file, line); + } else { + fprintf(stderr, "Error: Partial write at %s:%d. Expected %zu elements, wrote %zu\n", + file, line, nmemb, result); + } + fprintf(stderr, "Error details:\n"); + fprintf(stderr, " File: %s\n", file); + fprintf(stderr, " Line: %d\n", line); + fprintf(stderr, " Expected elements: %zu\n", nmemb); + fprintf(stderr, " Written elements: %zu\n", result); + exit(EXIT_FAILURE); + } +} + +#define fwriteCheck(ptr, size, nmemb, stream) fwrite_check(ptr, size, nmemb, stream, __FILE__, __LINE__) + // ---------------------------------------------------------------------------- // malloc error-handling wrapper util -void *malloc_check(size_t size, const char *file, int line) { +inline void *malloc_check(size_t size, const char *file, int line) { void *ptr = malloc(size); if (ptr == NULL) { fprintf(stderr, "Error: Memory allocation failed at %s:%d\n", file, line); @@ -108,7 +130,7 @@ void *malloc_check(size_t size, const char *file, int line) { // ---------------------------------------------------------------------------- // I/O ops -void create_dir_if_not_exists(const char *dir) { +inline void create_dir_if_not_exists(const char *dir) { if (dir == NULL) { return; } struct stat st = {0}; if (stat(dir, &st) == -1) { @@ -120,7 +142,7 @@ void create_dir_if_not_exists(const char *dir) { } } -int find_max_step(const char* output_log_dir) { +inline int find_max_step(const char* output_log_dir) { // find the DONE file in the log dir with highest step count if (output_log_dir == NULL) { return -1; } DIR* dir; diff --git a/test/device_file_io.cu b/test/device_file_io.cu new file mode 100644 index 0000000..d6e7f96 --- /dev/null +++ b/test/device_file_io.cu @@ -0,0 +1,53 @@ +#include "llmc/cuda_common.h" +#include +#include +#include + +void test(size_t nelem, size_t wt_buf, size_t rd_buf) { + + float* data; + cudaCheck(cudaMalloc(&data, nelem*sizeof(float))); + + // generate random array + std::vector random_data(nelem); + std::mt19937 rng(42); + std::uniform_real_distribution dist(-100.f, 100.f); + std::generate(random_data.begin(), random_data.end(), [&](){ return dist(rng); }); + + cudaCheck(cudaMemcpy(data, random_data.data(), random_data.size()*sizeof(float), cudaMemcpyHostToDevice)); + + cudaStream_t stream; + cudaStreamCreate(&stream); + + FILE* tmp = fopenCheck("tmp.bin", "w"); + device_to_file(tmp, data, nelem * sizeof(float), wt_buf, stream); + fcloseCheck(tmp); + + + float* reload; + cudaCheck(cudaMalloc(&reload, nelem*sizeof(float))); + + tmp = fopenCheck("tmp.bin", "r"); + file_to_device(reload, tmp, nelem * sizeof(float), rd_buf, stream); + fcloseCheck(tmp); + + std::vector cmp(nelem); + cudaCheck(cudaMemcpy(cmp.data(), reload, nelem * sizeof(float), cudaMemcpyDeviceToHost)); + for(int i = 0; i < nelem; ++i) { + if(random_data[i] != cmp[i]) { + fprintf(stderr, "FAIL: Mismatch at position %d: %f vs %f\n", i, random_data[i], cmp[i]); + exit(EXIT_FAILURE); + } + } + + cudaCheck(cudaFree(reload)); + cudaCheck(cudaFree(data)); +} + +int main() { + test(1025, 10000, 10000); // buffers larger than data + test(1025, 1024, 513); // different and smaller + test(500, 500*sizeof(float), + 500*sizeof(float)); // exact match + test(125'000, 10000, 10000); // large array +} \ No newline at end of file From 593a71c3743035065dd32e1d0f28c27b2009daed Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Thu, 6 Jun 2024 00:58:15 +0300 Subject: [PATCH 02/47] use new functions for checkpointing --- train_gpt2.cu | 47 +++++++++++++++++++++-------------------------- 1 file changed, 21 insertions(+), 26 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index e7b8718..68ceb23 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -383,10 +383,8 @@ void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { model_header[7] = model->config.padded_vocab_size; fwrite(model_header, sizeof(int), 256, model_file); // write the parameters - void* params_memory_cpu = (void*)mallocCheck(model->num_parameters_bytes); - cudaCheck(cudaMemcpy(params_memory_cpu, model->params_memory, model->num_parameters_bytes, cudaMemcpyDeviceToHost)); - fwrite(params_memory_cpu, 1, model->num_parameters_bytes, model_file); - free(params_memory_cpu); + device_to_file(model_file, model->params_memory, model->num_parameters_bytes, + 1024*1024*32, main_stream); // close file, we're done fcloseCheck(model_file); } @@ -449,10 +447,8 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { model->params_memory = malloc_and_point_parameters(&model->params, model->param_elements, model->param_sizeof); // read in all the parameters from file and copy them to device - void* params_memory_cpu = (void*)mallocCheck(model->num_parameters_bytes); - freadCheck(params_memory_cpu, 1, model->num_parameters_bytes, model_file); - cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, model->num_parameters_bytes, cudaMemcpyHostToDevice)); - free(params_memory_cpu); + file_to_device(model->params_memory, model_file, model->num_parameters_bytes, + 32*1024*1024, main_stream); fcloseCheck(model_file); // only return from this function once we are certain the params are ready on the GPU @@ -1183,16 +1179,11 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) // write AdamW m, v, and master_weights here (they are all float) size_t shard_num_parameters = multi_gpu_config.shard_num_parameters; - float* cpu_buffer = (float*)mallocCheck(shard_num_parameters * sizeof(float)); - cudaCheck(cudaMemcpy(cpu_buffer, model->m_memory, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); - fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); - cudaCheck(cudaMemcpy(cpu_buffer, model->v_memory, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); - fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); - if (model->use_master_weights == 1) { - cudaCheck(cudaMemcpy(cpu_buffer, model->master_weights, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); - fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + device_to_file(state_file, model->m_memory, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + device_to_file(state_file, model->v_memory, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + if(model->use_master_weights) { + device_to_file(state_file, model->master_weights, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); } - free(cpu_buffer); // write dataloader state if we are using the Permuted version of it if (loader->should_shuffle) { @@ -1231,20 +1222,24 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float))); } + + if(state_header[4] == 1 && !model->use_master_weights) { + printf0("Warning: Master weights are present in state, but not enabled for current run."); + } else if (state_header[4] == 0 && model->use_master_weights) { + printf0("Error: Master weights requested, but not present in state file."); + exit(EXIT_FAILURE); + } + if (model->master_weights == NULL && use_master_weights == 1) { printf0("allocating %zu MiB for master copy of params\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->master_weights, shard_num_parameters * sizeof(float))); } - float* cpu_buffer = (float*)mallocCheck(shard_num_parameters * sizeof(float)); - freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); - cudaCheck(cudaMemcpy(model->m_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); - freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); - cudaCheck(cudaMemcpy(model->v_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); - if (use_master_weights == 1) { - freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); - cudaCheck(cudaMemcpy(model->master_weights, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); + + file_to_device(model->m_memory, state_file, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + file_to_device(model->v_memory, state_file, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + if(model->use_master_weights) { + file_to_device(model->master_weights, state_file, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); } - free(cpu_buffer); // revive the DataLoader object and its state loader->should_shuffle = should_shuffle; From 188e7274e6b95cd1efab568685b8d0cb402f6aa5 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 01:22:31 +0300 Subject: [PATCH 03/47] add unit test to CI --- .github/workflows/ci_gpu.yml | 10 ++++++++++ {test => dev/test}/device_file_io.cu | 10 +++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) rename {test => dev/test}/device_file_io.cu (89%) diff --git a/.github/workflows/ci_gpu.yml b/.github/workflows/ci_gpu.yml index 9162d3b..ee61234 100644 --- a/.github/workflows/ci_gpu.yml +++ b/.github/workflows/ci_gpu.yml @@ -110,3 +110,13 @@ jobs: - name: Execute testing program fp32 with cuDNN run: ./test_gpt2fp32cu + + unit-tests-gpu: + runs-on: ubicloud-gpu-standard-1-latest + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Test Device<->File IO + run: nvcc -o device_file_io device_file_io.cu && ./device_file_io diff --git a/test/device_file_io.cu b/dev/test/device_file_io.cu similarity index 89% rename from test/device_file_io.cu rename to dev/test/device_file_io.cu index d6e7f96..fdb3e02 100644 --- a/test/device_file_io.cu +++ b/dev/test/device_file_io.cu @@ -1,4 +1,12 @@ -#include "llmc/cuda_common.h" +/* +Tests device <-> file IO functions + +compile and run as (from dev/test directory) +nvcc -o device_file_io device_file_io.cu && ./device_file_io +*/ + + +#include "../../llmc/cuda_common.h" #include #include #include From dbeb8fc551c6a21cc22f47a78f2fa81bee25b7e3 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 13:29:12 +0300 Subject: [PATCH 04/47] added missing checks --- test_gpt2.cu | 2 +- train_gpt2.cu | 18 +++++++++--------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/test_gpt2.cu b/test_gpt2.cu index 71f9d07..5912362 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -166,7 +166,7 @@ int main(int argc, char *argv[]) { // copy logits to CPU so we can compare them floatX* logits_cpu_raw = (floatX*)mallocCheck(B * T * Vp * sizeof(floatX)); float* logits_cpu = (float*)mallocCheck(B * T * Vp * sizeof(float)); - cudaMemcpy(logits_cpu_raw, model.acts.output, B * T * Vp * sizeof(floatX), cudaMemcpyDeviceToHost); + cudaCheck(cudaMemcpy(logits_cpu_raw, model.acts.output, B * T * Vp * sizeof(floatX), cudaMemcpyDeviceToHost)); for (int i = 0; i < B * T * Vp; i++) { logits_cpu[i] = (float)logits_cpu_raw[i]; } diff --git a/train_gpt2.cu b/train_gpt2.cu index 68ceb23..3d3e439 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -381,7 +381,7 @@ void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { model_header[5] = model->config.num_heads; model_header[6] = model->config.channels; model_header[7] = model->config.padded_vocab_size; - fwrite(model_header, sizeof(int), 256, model_file); + fwriteCheck(model_header, sizeof(int), 256, model_file); // write the parameters device_to_file(model_file, model->params_memory, model->num_parameters_bytes, 1024*1024*32, main_stream); @@ -1175,7 +1175,7 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) // 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 - fwrite(state_header, sizeof(int), 256, state_file); + fwriteCheck(state_header, sizeof(int), 256, state_file); // write AdamW m, v, and master_weights here (they are all float) size_t shard_num_parameters = multi_gpu_config.shard_num_parameters; @@ -1187,13 +1187,13 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) // write dataloader state if we are using the Permuted version of it if (loader->should_shuffle) { - fwrite(&loader->glob_result.gl_pathc, sizeof(size_t), 1, state_file); // number of shards - fwrite(loader->shard_indices, sizeof(int), loader->glob_result.gl_pathc, state_file); - fwrite(&loader->shard_num_samples, sizeof(size_t), 1, state_file); - fwrite(loader->intra_shard_indices, sizeof(int), loader->shard_num_samples, state_file); - fwrite(&loader->shuffle_rng, sizeof(mt19937_state), 1, state_file); + fwriteCheck(&loader->glob_result.gl_pathc, sizeof(size_t), 1, state_file); // number of shards + fwriteCheck(loader->shard_indices, sizeof(int), loader->glob_result.gl_pathc, state_file); + fwriteCheck(&loader->shard_num_samples, sizeof(size_t), 1, state_file); + fwriteCheck(loader->intra_shard_indices, sizeof(int), loader->shard_num_samples, state_file); + fwriteCheck(&loader->shuffle_rng, sizeof(mt19937_state), 1, state_file); } - fclose(state_file); + fcloseCheck(state_file); } void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename) { @@ -1264,7 +1264,7 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename dataloader_resume(loader, current_shard_idx, current_sample_idx); // all done, close state file - fclose(state_file); + fcloseCheck(state_file); } From 33136e0b082ffebe2c05b7c9c48260287080c0f1 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 14:16:59 +0300 Subject: [PATCH 05/47] fix compilation with clang --- llmc/utils.h | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/llmc/utils.h b/llmc/utils.h index e533f2b..74a06d2 100644 --- a/llmc/utils.h +++ b/llmc/utils.h @@ -21,7 +21,7 @@ // simple replace fopen, fread, fclose, fseek // with fopenCheck, freadCheck, fcloseCheck, fseekCheck -inline FILE *fopen_check(const char *path, const char *mode, const char *file, int line) { +extern inline FILE *fopen_check(const char *path, const char *mode, const char *file, int line) { FILE *fp = fopen(path, mode); if (fp == NULL) { fprintf(stderr, "Error: Failed to open file '%s' at %s:%d\n", path, file, line); @@ -39,7 +39,7 @@ inline FILE *fopen_check(const char *path, const char *mode, const char *file, i #define fopenCheck(path, mode) fopen_check(path, mode, __FILE__, __LINE__) -inline void fread_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { +extern inline void fread_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { size_t result = fread(ptr, size, nmemb, stream); if (result != nmemb) { if (feof(stream)) { @@ -61,7 +61,7 @@ inline void fread_check(void *ptr, size_t size, size_t nmemb, FILE *stream, cons #define freadCheck(ptr, size, nmemb, stream) fread_check(ptr, size, nmemb, stream, __FILE__, __LINE__) -inline void fclose_check(FILE *fp, const char *file, int line) { +extern inline void fclose_check(FILE *fp, const char *file, int line) { if (fclose(fp) != 0) { fprintf(stderr, "Error: Failed to close file at %s:%d\n", file, line); fprintf(stderr, "Error details:\n"); @@ -73,7 +73,7 @@ inline void fclose_check(FILE *fp, const char *file, int line) { #define fcloseCheck(fp) fclose_check(fp, __FILE__, __LINE__) -inline void fseek_check(FILE *fp, long off, int whence, const char *file, int line) { +extern inline void fseek_check(FILE *fp, long off, int whence, const char *file, int line) { if (fseek(fp, off, whence) != 0) { fprintf(stderr, "Error: Failed to seek in file at %s:%d\n", file, line); fprintf(stderr, "Error details:\n"); @@ -87,7 +87,7 @@ inline void fseek_check(FILE *fp, long off, int whence, const char *file, int li #define fseekCheck(fp, off, whence) fseek_check(fp, off, whence, __FILE__, __LINE__) -inline void fwrite_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { +extern inline void fwrite_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { size_t result = fwrite(ptr, size, nmemb, stream); if (result != nmemb) { if (feof(stream)) { @@ -112,7 +112,7 @@ inline void fwrite_check(void *ptr, size_t size, size_t nmemb, FILE *stream, con // ---------------------------------------------------------------------------- // malloc error-handling wrapper util -inline void *malloc_check(size_t size, const char *file, int line) { +extern inline void *malloc_check(size_t size, const char *file, int line) { void *ptr = malloc(size); if (ptr == NULL) { fprintf(stderr, "Error: Memory allocation failed at %s:%d\n", file, line); @@ -130,7 +130,7 @@ inline void *malloc_check(size_t size, const char *file, int line) { // ---------------------------------------------------------------------------- // I/O ops -inline void create_dir_if_not_exists(const char *dir) { +extern inline void create_dir_if_not_exists(const char *dir) { if (dir == NULL) { return; } struct stat st = {0}; if (stat(dir, &st) == -1) { @@ -142,7 +142,7 @@ inline void create_dir_if_not_exists(const char *dir) { } } -inline int find_max_step(const char* output_log_dir) { +extern inline int find_max_step(const char* output_log_dir) { // find the DONE file in the log dir with highest step count if (output_log_dir == NULL) { return -1; } DIR* dir; From 37c3815ede7c604a87570bbb542db6450afe27a1 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 14:35:05 +0300 Subject: [PATCH 06/47] fix path --- .github/workflows/ci_gpu.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci_gpu.yml b/.github/workflows/ci_gpu.yml index ee61234..ac2e1f4 100644 --- a/.github/workflows/ci_gpu.yml +++ b/.github/workflows/ci_gpu.yml @@ -119,4 +119,4 @@ jobs: uses: actions/checkout@v4 - name: Test Device<->File IO - run: nvcc -o device_file_io device_file_io.cu && ./device_file_io + run: cd dev/test && nvcc -o device_file_io device_file_io.cu && ./device_file_io From 3ebd5abe82d1394fa807932cad2c37d57b58311d Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 19:58:03 +0300 Subject: [PATCH 07/47] fixup tests --- dev/test/device_file_io.cu | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/dev/test/device_file_io.cu b/dev/test/device_file_io.cu index fdb3e02..71fb1ce 100644 --- a/dev/test/device_file_io.cu +++ b/dev/test/device_file_io.cu @@ -9,9 +9,10 @@ nvcc -o device_file_io device_file_io.cu && ./device_file_io #include "../../llmc/cuda_common.h" #include #include +#include #include -void test(size_t nelem, size_t wt_buf, size_t rd_buf) { +void test(size_t nelem, size_t wt_buf_size, size_t rd_buf_size) { float* data; cudaCheck(cudaMalloc(&data, nelem*sizeof(float))); @@ -28,7 +29,7 @@ void test(size_t nelem, size_t wt_buf, size_t rd_buf) { cudaStreamCreate(&stream); FILE* tmp = fopenCheck("tmp.bin", "w"); - device_to_file(tmp, data, nelem * sizeof(float), wt_buf, stream); + device_to_file(tmp, data, nelem * sizeof(float), wt_buf_size, stream); fcloseCheck(tmp); @@ -36,20 +37,22 @@ void test(size_t nelem, size_t wt_buf, size_t rd_buf) { cudaCheck(cudaMalloc(&reload, nelem*sizeof(float))); tmp = fopenCheck("tmp.bin", "r"); - file_to_device(reload, tmp, nelem * sizeof(float), rd_buf, stream); + file_to_device(reload, tmp, nelem * sizeof(float), rd_buf_size, stream); fcloseCheck(tmp); std::vector cmp(nelem); cudaCheck(cudaMemcpy(cmp.data(), reload, nelem * sizeof(float), cudaMemcpyDeviceToHost)); for(int i = 0; i < nelem; ++i) { - if(random_data[i] != cmp[i]) { + if(random_data[i] != cmp[i]) { fprintf(stderr, "FAIL: Mismatch at position %d: %f vs %f\n", i, random_data[i], cmp[i]); + remove("tmp.bin"); exit(EXIT_FAILURE); } } cudaCheck(cudaFree(reload)); cudaCheck(cudaFree(data)); + remove("tmp.bin"); } int main() { From ec4424013a6a5e9c91b29db29237a4438cc568cb Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 20:00:25 +0300 Subject: [PATCH 08/47] made buffer size more easily configurable --- train_gpt2.cu | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 3d3e439..f461478 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -71,6 +71,8 @@ cudaDeviceProp deviceProp; // fills in common_start() cudaStream_t main_stream; // one global variable to hold the multi-GPU configuration for this process MultiGpuConfig multi_gpu_config; +// buffer size to use for device <-> disk io +constexpr const size_t IO_BUF_SIZE = 32 * 1024 * 1024; // convenience function that only prints if the rank of process is zero void printf0(const char *format, ...) { @@ -448,7 +450,7 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { // read in all the parameters from file and copy them to device file_to_device(model->params_memory, model_file, model->num_parameters_bytes, - 32*1024*1024, main_stream); + IO_BUF_SIZE, main_stream); fcloseCheck(model_file); // only return from this function once we are certain the params are ready on the GPU @@ -1179,10 +1181,10 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) // write AdamW m, v, and master_weights here (they are all float) size_t shard_num_parameters = multi_gpu_config.shard_num_parameters; - device_to_file(state_file, model->m_memory, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); - device_to_file(state_file, model->v_memory, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + device_to_file(state_file, model->m_memory, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); + device_to_file(state_file, model->v_memory, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); if(model->use_master_weights) { - device_to_file(state_file, model->master_weights, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + device_to_file(state_file, model->master_weights, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); } // write dataloader state if we are using the Permuted version of it @@ -1235,10 +1237,10 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename cudaCheck(cudaMalloc((void**)&model->master_weights, shard_num_parameters * sizeof(float))); } - file_to_device(model->m_memory, state_file, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); - file_to_device(model->v_memory, state_file, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + file_to_device(model->m_memory, state_file, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); + file_to_device(model->v_memory, state_file, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); if(model->use_master_weights) { - file_to_device(model->master_weights, state_file, shard_num_parameters * sizeof(float), 32*1024*1024, main_stream); + file_to_device(model->master_weights, state_file, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); } // revive the DataLoader object and its state From d95313d0a9e72ad4a1745300641d1694afbb7842 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 20:02:07 +0300 Subject: [PATCH 09/47] explicit sync --- llmc/cuda_common.h | 1 + 1 file changed, 1 insertion(+) diff --git a/llmc/cuda_common.h b/llmc/cuda_common.h index 921bf23..2d973e4 100644 --- a/llmc/cuda_common.h +++ b/llmc/cuda_common.h @@ -194,6 +194,7 @@ inline void file_to_device(void* dest, FILE* src, size_t num_bytes, size_t buffe // copy the last remaining write buffer to gpu cudaCheck(cudaMemcpyAsync(gpu_write_ptr, write_buffer, write_buffer_size, cudaMemcpyHostToDevice, stream)); + cudaCheck(cudaStreamSynchronize(stream)); cudaCheck(cudaFreeHost(buffer_space)); } From b3b4dabab08bd513546a8f491adef9d04528cac7 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Tue, 18 Jun 2024 20:03:28 +0300 Subject: [PATCH 10/47] small touchups --- train_gpt2.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index f461478..2a9271e 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1225,9 +1225,9 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float))); } - if(state_header[4] == 1 && !model->use_master_weights) { + if(use_master_weights == 1 && !model->use_master_weights) { printf0("Warning: Master weights are present in state, but not enabled for current run."); - } else if (state_header[4] == 0 && model->use_master_weights) { + } else if (use_master_weights == 0 && model->use_master_weights) { printf0("Error: Master weights requested, but not present in state file."); exit(EXIT_FAILURE); } From 66950f37907c6701ffbf5f4477a2ac389688da89 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Wed, 19 Jun 2024 17:45:34 +0200 Subject: [PATCH 11/47] Cast computation to uint --- llmc/cuda_utils.cuh | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index fe3b265..3205938 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -183,8 +183,11 @@ __device__ __host__ constexpr unsigned int SquirrelNoise5(int positionX, unsigne } __device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY, unsigned int seed) { - constexpr int PRIME_NUMBER = 198491317; // Large prime number with non-boring bits - return SquirrelNoise5(indexX + (PRIME_NUMBER * indexY), seed); + constexpr unsigned int PRIME_NUMBER = 198491317u; // Large prime number with non-boring bits + unsigned int x = static_cast(indexX); + unsigned int y = static_cast(indexY); + + return SquirrelNoise5(x + (PRIME_NUMBER * y), seed); } // stochastic rounding built on top of Squirel Noise above (with seed updated per step via xorshift) From bb9eab87360941abaa11bbc033af74dc13b74193 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Wed, 19 Jun 2024 21:05:28 +0200 Subject: [PATCH 12/47] Change positionX to uint --- llmc/cuda_utils.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 3205938..83edc2c 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -160,14 +160,14 @@ __device__ inline float blockReduce(float val, bool final_sync=false, float out_ // This gives us a random number from threadIdx/blockIdx + a single seed for the entire GPU // todo - possibly overkill and we don't need such high quality random numbers? (tbd) // http://eiserloh.net/noise/SquirrelNoise5.hpp -__device__ __host__ constexpr unsigned int SquirrelNoise5(int positionX, unsigned int seed) +__device__ __host__ constexpr unsigned int SquirrelNoise5(unsigned int positionX, unsigned int seed) { constexpr unsigned int SQ5_BIT_NOISE1 = 0xd2a80a3f; // 11010010101010000000101000111111 constexpr unsigned int SQ5_BIT_NOISE2 = 0xa884f197; // 10101000100001001111000110010111 constexpr unsigned int SQ5_BIT_NOISE3 = 0x6C736F4B; // 01101100011100110110111101001011 constexpr unsigned int SQ5_BIT_NOISE4 = 0xB79F3ABB; // 10110111100111110011101010111011 constexpr unsigned int SQ5_BIT_NOISE5 = 0x1b56c4f5; // 00011011010101101100010011110101 - unsigned int mangledBits = (unsigned int) positionX; + unsigned int mangledBits = positionX; mangledBits *= SQ5_BIT_NOISE1; mangledBits += seed; mangledBits ^= (mangledBits >> 9); From 8249736233264e798d4f6dadbdf656b23a04db9c Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 16:36:17 +0200 Subject: [PATCH 13/47] Add learning rate schedulers --- llmc/schedulers.h | 43 +++++++++++++++++++++++++++++++++++++++++++ train_gpt2.cu | 20 ++++++++------------ 2 files changed, 51 insertions(+), 12 deletions(-) create mode 100644 llmc/schedulers.h diff --git a/llmc/schedulers.h b/llmc/schedulers.h new file mode 100644 index 0000000..bf469aa --- /dev/null +++ b/llmc/schedulers.h @@ -0,0 +1,43 @@ + + +// Cosine learning rate scheduler + +#ifndef SCHEDULERS_H + +#define SCHEDULERS_H + +#include +#include + +typedef struct { + float learning_rate; + int warmup_iterations; + int train_num_batches; + float final_learning_rate_frac; +} CosineLearningRateScheduler; + + +// learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac +float get_learning_rate(CosineLearningRateScheduler *scheduler, int step) { + float step_learning_rate = scheduler->learning_rate; + if (step < scheduler->warmup_iterations) { + step_learning_rate = scheduler->learning_rate * ((float)(step + 1)) / scheduler->warmup_iterations; + } else { + float decay_ratio = ((float)(step - scheduler->warmup_iterations)) / (scheduler->train_num_batches - scheduler->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 = scheduler->learning_rate * scheduler->final_learning_rate_frac; + step_learning_rate = min_lr + coeff * (scheduler->learning_rate - min_lr); + } + return step_learning_rate; +} + +void lr_scheduler_init(CosineLearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { + scheduler->learning_rate = learning_rate; + scheduler->warmup_iterations = warmup_iterations; + scheduler->train_num_batches = train_num_batches; + scheduler->final_learning_rate_frac = final_learning_rate_frac; +} + +#endif // SCHEDULERS_H \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index f9e846a..e0dac38 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -21,6 +21,8 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. #include "llmc/dataloader.h" // defines: manual_seed, normal_ (same as torch.manual_seed and torch.normal) #include "llmc/rand.h" +// defines learning rate schedulers +#include "llmc/schedulers.h" // defines: sample_softmax, random_f32 #include "llmc/sampler.h" // defines: logger_init, logger_log_eval, logger_log_val, logger_log_train @@ -1546,6 +1548,10 @@ int main(int argc, char *argv[]) { Tokenizer tokenizer; tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); + // set up learning rate scheduler + CosineLearningRateScheduler lr_scheduler; + lr_scheduler_init(&lr_scheduler, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); + // some memory for generating samples from the model int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); floatX* cpu_logits_raw = (floatX*)mallocCheck(model.config.vocab_size * sizeof(floatX)); @@ -1704,18 +1710,8 @@ int main(int argc, char *argv[]) { model.mean_loss = lossf; // average the loss and the gradients between all processes gpt2_multi_gpu_loss_reduce(&model, &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); - } + // fetch the next learning rate + float step_learning_rate = get_learning_rate(&lr_scheduler, step); // update the model parameters float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, 1.0f, step+1, &multi_gpu_config); // zero out the gradients for the next iteration From 2eaea3b397f2649956f6cd30ec20e9530fc50923 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 16:40:07 +0200 Subject: [PATCH 14/47] Add schedulers header --- llmc/schedulers.h | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index bf469aa..bd0f5f4 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -1,7 +1,6 @@ - - -// Cosine learning rate scheduler - +/* +Implements various learning rate schedulers. +*/ #ifndef SCHEDULERS_H #define SCHEDULERS_H @@ -16,7 +15,6 @@ typedef struct { float final_learning_rate_frac; } CosineLearningRateScheduler; - // learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac float get_learning_rate(CosineLearningRateScheduler *scheduler, int step) { float step_learning_rate = scheduler->learning_rate; From 4a1c5a9c0f897f16b50827b6fc8f9e4848e686ea Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 17:11:27 +0200 Subject: [PATCH 15/47] Add cyclic triangular LR scheduler --- llmc/schedulers.h | 46 +++++++++++++++++++++++++++++++++++++++------- train_gpt2.cu | 4 ++-- 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index bd0f5f4..e1d3c71 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -8,6 +8,10 @@ Implements various learning rate schedulers. #include #include +// +// Learning rate scheduler structs +// + typedef struct { float learning_rate; int warmup_iterations; @@ -15,27 +19,55 @@ typedef struct { float final_learning_rate_frac; } CosineLearningRateScheduler; -// learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac -float get_learning_rate(CosineLearningRateScheduler *scheduler, int step) { - float step_learning_rate = scheduler->learning_rate; +typedef struct { + float min_lr; + float max_lr; + int step_size; +} CyclicTriangularLearningRateScheduler; + +// +// Learning rate scheduler functions +// + +// cosine learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac +float get_learning_rate_cosine(CosineLearningRateScheduler *scheduler, int step) { + float lr = scheduler->learning_rate; if (step < scheduler->warmup_iterations) { - step_learning_rate = scheduler->learning_rate * ((float)(step + 1)) / scheduler->warmup_iterations; + lr = scheduler->learning_rate * ((float)(step + 1)) / scheduler->warmup_iterations; } else { float decay_ratio = ((float)(step - scheduler->warmup_iterations)) / (scheduler->train_num_batches - scheduler->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 = scheduler->learning_rate * scheduler->final_learning_rate_frac; - step_learning_rate = min_lr + coeff * (scheduler->learning_rate - min_lr); + lr = min_lr + coeff * (scheduler->learning_rate - min_lr); } - return step_learning_rate; + return lr; } -void lr_scheduler_init(CosineLearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { +// cyclic triangular learning rate schedule: linearly increase LR from min LR to max LR, then linearly decrease LR to min LR (repeat) +float get_learning_rate_triangular(CyclicTriangularLearningRateScheduler *scheduler, int step) { + int cycle = 1 + step / (2 * scheduler->step_size); + float x = fabsf((float)step / scheduler->step_size - 2 * cycle + 1); + float lr = scheduler->min_lr + (scheduler->max_lr - scheduler->min_lr) * fmaxf(0, (1 - x)); + return lr; +} + +// +// Init functions +// + +void lr_scheduler_init_cosine(CosineLearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { scheduler->learning_rate = learning_rate; scheduler->warmup_iterations = warmup_iterations; scheduler->train_num_batches = train_num_batches; scheduler->final_learning_rate_frac = final_learning_rate_frac; } +void lr_scheduler_init_triangular(CyclicTriangularLearningRateScheduler *scheduler, float min_lr, float max_lr, int step_size) { + scheduler->min_lr = min_lr; + scheduler->max_lr = max_lr; + scheduler->step_size = step_size; +} + #endif // SCHEDULERS_H \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index e0dac38..24984f3 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1550,7 +1550,7 @@ int main(int argc, char *argv[]) { // set up learning rate scheduler CosineLearningRateScheduler lr_scheduler; - lr_scheduler_init(&lr_scheduler, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); + lr_scheduler_init_cosine(&lr_scheduler, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); // some memory for generating samples from the model int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); @@ -1711,7 +1711,7 @@ int main(int argc, char *argv[]) { // average the loss and the gradients between all processes gpt2_multi_gpu_loss_reduce(&model, &multi_gpu_config); // fetch the next learning rate - float step_learning_rate = get_learning_rate(&lr_scheduler, step); + float step_learning_rate = get_learning_rate_cosine(&lr_scheduler, step); // update the model parameters float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, 1.0f, step+1, &multi_gpu_config); // zero out the gradients for the next iteration From 6f3d099e88870b5ea72b24c2bf63f9232481cca5 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 17:16:59 +0200 Subject: [PATCH 16/47] Add constant lr scheduler --- llmc/schedulers.h | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index e1d3c71..0d6783a 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -25,6 +25,11 @@ typedef struct { int step_size; } CyclicTriangularLearningRateScheduler; +// Constant learning rate scheduler +typedef struct { + float learning_rate; +} ConstantLearningRateScheduler; + // // Learning rate scheduler functions // @@ -47,12 +52,17 @@ float get_learning_rate_cosine(CosineLearningRateScheduler *scheduler, int step) // cyclic triangular learning rate schedule: linearly increase LR from min LR to max LR, then linearly decrease LR to min LR (repeat) float get_learning_rate_triangular(CyclicTriangularLearningRateScheduler *scheduler, int step) { - int cycle = 1 + step / (2 * scheduler->step_size); - float x = fabsf((float)step / scheduler->step_size - 2 * cycle + 1); + int cycle_index = 1 + step / (2 * scheduler->step_size); // tells us which cycle we are in, starting at 1 + float x = fabsf((float)step / scheduler->step_size - 2 * cycle_index + 1); // goes from 0 to 1 to 0 float lr = scheduler->min_lr + (scheduler->max_lr - scheduler->min_lr) * fmaxf(0, (1 - x)); return lr; } +// constant learning rate schedule +float get_learning_rate_constant(ConstantLearningRateScheduler *scheduler, int step) { + return scheduler->learning_rate; +} + // // Init functions // @@ -70,4 +80,8 @@ void lr_scheduler_init_triangular(CyclicTriangularLearningRateScheduler *schedul scheduler->step_size = step_size; } +void lr_scheduler_init_constant(ConstantLearningRateScheduler *scheduler, float learning_rate) { + scheduler->learning_rate = learning_rate; +} + #endif // SCHEDULERS_H \ No newline at end of file From 7308c2831cef2ad965ee8a7e8f841386fe4f567b Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 17 Jun 2024 17:27:24 +0200 Subject: [PATCH 17/47] Add linear LR scheduler --- llmc/schedulers.h | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index 0d6783a..74d703f 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -19,6 +19,14 @@ typedef struct { float final_learning_rate_frac; } CosineLearningRateScheduler; +// Linear with warmup learning rate scheduler +typedef struct { + float learning_rate; + int warmup_iterations; + int train_num_batches; + float final_learning_rate_frac; +} LinearLearningRateScheduler; + typedef struct { float min_lr; float max_lr; @@ -50,6 +58,20 @@ float get_learning_rate_cosine(CosineLearningRateScheduler *scheduler, int step) return lr; } +// linear warmup learning rate schedule: warmup linearly to max LR, then decay linearly to LR * final_learning_rate_frac +float get_learning_rate_linear(LinearLearningRateScheduler *scheduler, int step) { + float lr = scheduler->learning_rate; + if (step < scheduler->warmup_iterations) { + lr = scheduler->learning_rate * ((float)(step + 1)) / scheduler->warmup_iterations; + } else { + float decay_ratio = ((float)(step - scheduler->warmup_iterations)) / (scheduler->train_num_batches - scheduler->warmup_iterations); + assert(0.0f <= decay_ratio && decay_ratio <= 1.0f); + float min_lr = scheduler->learning_rate * scheduler->final_learning_rate_frac; + lr = scheduler->learning_rate - decay_ratio * (scheduler->learning_rate - min_lr); + } + return lr; +} + // cyclic triangular learning rate schedule: linearly increase LR from min LR to max LR, then linearly decrease LR to min LR (repeat) float get_learning_rate_triangular(CyclicTriangularLearningRateScheduler *scheduler, int step) { int cycle_index = 1 + step / (2 * scheduler->step_size); // tells us which cycle we are in, starting at 1 @@ -74,6 +96,13 @@ void lr_scheduler_init_cosine(CosineLearningRateScheduler *scheduler, float lear scheduler->final_learning_rate_frac = final_learning_rate_frac; } +void lr_scheduler_init_linear(LinearLearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { + scheduler->learning_rate = learning_rate; + scheduler->warmup_iterations = warmup_iterations; + scheduler->train_num_batches = train_num_batches; + scheduler->final_learning_rate_frac = final_learning_rate_frac; +} + void lr_scheduler_init_triangular(CyclicTriangularLearningRateScheduler *scheduler, float min_lr, float max_lr, int step_size) { scheduler->min_lr = min_lr; scheduler->max_lr = max_lr; From 9d141ddab9c7fc6e1adf42a43769ba035a9e31ce Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Wed, 19 Jun 2024 18:36:45 +0200 Subject: [PATCH 18/47] Add ugly switching between schedulers --- llmc/schedulers.h | 32 ++++++++++++++++++++++++++++++++ train_gpt2.cu | 40 +++++++++++++++++++++++++++++++++++++--- 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index 74d703f..b73be54 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -8,6 +8,38 @@ Implements various learning rate schedulers. #include #include +typedef enum { + LR_SCHEDULER_COSINE, + LR_SCHEDULER_LINEAR, + LR_SCHEDULER_TRIANGULAR, + LR_SCHEDULER_CONSTANT, + NUM_LR_SCHEDULERS // To keep track of the number of schedulers +} LRSchedulerType; + +const char* lr_scheduler_names[] = { + "cosine", + "linear", + "triangular", + "constant", +}; + +const char* get_lr_scheduler_name(LRSchedulerType type) { + if (type < 0 || type >= NUM_LR_SCHEDULERS) { + exit(EXIT_FAILURE); + } + return lr_scheduler_names[type]; +} + +LRSchedulerType get_lr_scheduler_type_from_name(const char* name) { + for (int i = 0; i < NUM_LR_SCHEDULERS; ++i) { + if (strcmp(name, lr_scheduler_names[i]) == 0) { + return (LRSchedulerType)i; + } + } + printf("Warning: Unknown learning rate scheduler name: %s\n. Using cosine as default.", name); + return LR_SCHEDULER_COSINE; // Default to cosine if not found +} + // // Learning rate scheduler structs // diff --git a/train_gpt2.cu b/train_gpt2.cu index 24984f3..d311462 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1303,6 +1303,7 @@ void error_usage() { // workload (number of steps) fprintf(stderr, " -x max_steps of optimization to run (-1 (default) = disable, run 1 epoch)\n"); // optimization + fprintf(stderr, " -k learning rate scheduler (default = cosine)\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"); @@ -1333,6 +1334,7 @@ int main(int argc, char *argv[]) { const char* train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; const char* val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; const char* load_filename = "gpt2_124M_bf16.bin"; // bf16 weights of the model + LRSchedulerType lr_scheduler_type = LR_SCHEDULER_COSINE; const char* output_log_dir = NULL; int checkpoint_every = 0; // write optimization checkpoints every how many steps? int resume = 0; // resume the optimization, if one is found inside output_log_dir? @@ -1383,6 +1385,7 @@ int main(int argc, char *argv[]) { 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 if (argv[i][1] == 'k') { lr_scheduler_type = get_lr_scheduler_type_from_name(argv[i + 1]); } else { error_usage(); } } // should do a bit more error checking here @@ -1416,6 +1419,7 @@ 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("| LR scheduler | %-50s |\n", get_lr_scheduler_name(lr_scheduler_type)); 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); @@ -1549,8 +1553,26 @@ int main(int argc, char *argv[]) { tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); // set up learning rate scheduler - CosineLearningRateScheduler lr_scheduler; - lr_scheduler_init_cosine(&lr_scheduler, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); + CosineLearningRateScheduler lr_scheduler_cosine; + LinearLearningRateScheduler lr_scheduler_linear; + CyclicTriangularLearningRateScheduler lr_scheduler_triangular; + ConstantLearningRateScheduler lr_scheduler_constant; + if (lr_scheduler_type == LR_SCHEDULER_COSINE) { + lr_scheduler_init_cosine(&lr_scheduler_cosine, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); + } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { + lr_scheduler_init_linear(&lr_scheduler_linear, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); + } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { + // Hardcode some reasonable defaults for now + float min_lr = learning_rate / 10.0f; + float max_lr = learning_rate; + int step_size = train_num_batches / 4; + lr_scheduler_init_triangular(&lr_scheduler_triangular, min_lr, max_lr, step_size); + } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { + lr_scheduler_init_constant(&lr_scheduler_constant, learning_rate); + } else { + printf("Unknown learning rate scheduler type\n"); + exit(EXIT_FAILURE); + } // some memory for generating samples from the model int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); @@ -1711,7 +1733,19 @@ int main(int argc, char *argv[]) { // average the loss and the gradients between all processes gpt2_multi_gpu_loss_reduce(&model, &multi_gpu_config); // fetch the next learning rate - float step_learning_rate = get_learning_rate_cosine(&lr_scheduler, step); + float step_learning_rate; + if (lr_scheduler_type == LR_SCHEDULER_COSINE) { + step_learning_rate = get_learning_rate_cosine(&lr_scheduler_cosine, step); + } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { + step_learning_rate = get_learning_rate_linear(&lr_scheduler_linear, step); + } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { + step_learning_rate = get_learning_rate_triangular(&lr_scheduler_triangular, step); + } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { + step_learning_rate = get_learning_rate_constant(&lr_scheduler_constant, step); + } else { + printf("Unknown learning rate scheduler type\n"); + exit(EXIT_FAILURE); + } // update the model parameters float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, 1.0f, step+1, &multi_gpu_config); // zero out the gradients for the next iteration From 7416ee425fb23da3700bf9c4449b9ef930cb74ff Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Wed, 19 Jun 2024 20:39:17 +0200 Subject: [PATCH 19/47] Simplify logic since they all share similar args --- llmc/schedulers.h | 93 +++++++++++++++++++---------------------------- train_gpt2.cu | 38 ++----------------- 2 files changed, 42 insertions(+), 89 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index b73be54..b82dbdf 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -41,7 +41,7 @@ LRSchedulerType get_lr_scheduler_type_from_name(const char* name) { } // -// Learning rate scheduler structs +// Learning rate scheduler structs and init // typedef struct { @@ -49,33 +49,39 @@ typedef struct { int warmup_iterations; int train_num_batches; float final_learning_rate_frac; -} CosineLearningRateScheduler; +} LearningRateScheduler; -// Linear with warmup learning rate scheduler -typedef struct { - float learning_rate; - int warmup_iterations; - int train_num_batches; - float final_learning_rate_frac; -} LinearLearningRateScheduler; - -typedef struct { - float min_lr; - float max_lr; - int step_size; -} CyclicTriangularLearningRateScheduler; - -// Constant learning rate scheduler -typedef struct { - float learning_rate; -} ConstantLearningRateScheduler; +void lr_scheduler_init(LearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { + scheduler->learning_rate = learning_rate; + scheduler->warmup_iterations = warmup_iterations; + scheduler->train_num_batches = train_num_batches; + scheduler->final_learning_rate_frac = final_learning_rate_frac; +} // // Learning rate scheduler functions // +// switch to the appropriate learning rate scheduler +float get_learning_rate(LRSchedulerType lr_scheduler_type, LearningRateScheduler *scheduler, int step) { + float step_learning_rate; + if (lr_scheduler_type == LR_SCHEDULER_COSINE) { + step_learning_rate = get_learning_rate_cosine(scheduler, step); + } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { + step_learning_rate = get_learning_rate_linear(scheduler, step); + } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { + step_learning_rate = get_learning_rate_triangular(scheduler, step); + } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { + step_learning_rate = get_learning_rate_constant(scheduler, step); + } else { + printf("Unknown learning rate scheduler type\n"); + exit(EXIT_FAILURE); + } + return step_learning_rate; +} + // cosine learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac -float get_learning_rate_cosine(CosineLearningRateScheduler *scheduler, int step) { +float get_learning_rate_cosine(LearningRateScheduler *scheduler, int step) { float lr = scheduler->learning_rate; if (step < scheduler->warmup_iterations) { lr = scheduler->learning_rate * ((float)(step + 1)) / scheduler->warmup_iterations; @@ -91,7 +97,7 @@ float get_learning_rate_cosine(CosineLearningRateScheduler *scheduler, int step) } // linear warmup learning rate schedule: warmup linearly to max LR, then decay linearly to LR * final_learning_rate_frac -float get_learning_rate_linear(LinearLearningRateScheduler *scheduler, int step) { +float get_learning_rate_linear(LearningRateScheduler *scheduler, int step) { float lr = scheduler->learning_rate; if (step < scheduler->warmup_iterations) { lr = scheduler->learning_rate * ((float)(step + 1)) / scheduler->warmup_iterations; @@ -105,44 +111,21 @@ float get_learning_rate_linear(LinearLearningRateScheduler *scheduler, int step) } // cyclic triangular learning rate schedule: linearly increase LR from min LR to max LR, then linearly decrease LR to min LR (repeat) -float get_learning_rate_triangular(CyclicTriangularLearningRateScheduler *scheduler, int step) { - int cycle_index = 1 + step / (2 * scheduler->step_size); // tells us which cycle we are in, starting at 1 - float x = fabsf((float)step / scheduler->step_size - 2 * cycle_index + 1); // goes from 0 to 1 to 0 - float lr = scheduler->min_lr + (scheduler->max_lr - scheduler->min_lr) * fmaxf(0, (1 - x)); +// currently hardcoded to support only a single cycle +float get_learning_rate_triangular(LearningRateScheduler *scheduler, int step) { + int step_size = scheduler->train_num_batches / 2; // number of steps in half a cycle + float min_lr = scheduler->learning_rate * scheduler->final_learning_rate_frac; + float max_lr = scheduler->learning_rate; + + int cycle_index = 1 + step / (2 * step_size); // tells us which cycle we are in, starting at 1 + float x = fabsf((float)step / step_size - 2 * cycle_index + 1); // goes from 0 to 1 to 0 + float lr = min_lr + (max_lr - min_lr) * fmaxf(0, (1 - x)); return lr; } // constant learning rate schedule -float get_learning_rate_constant(ConstantLearningRateScheduler *scheduler, int step) { +float get_learning_rate_constant(LearningRateScheduler *scheduler, int step) { return scheduler->learning_rate; } -// -// Init functions -// - -void lr_scheduler_init_cosine(CosineLearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { - scheduler->learning_rate = learning_rate; - scheduler->warmup_iterations = warmup_iterations; - scheduler->train_num_batches = train_num_batches; - scheduler->final_learning_rate_frac = final_learning_rate_frac; -} - -void lr_scheduler_init_linear(LinearLearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { - scheduler->learning_rate = learning_rate; - scheduler->warmup_iterations = warmup_iterations; - scheduler->train_num_batches = train_num_batches; - scheduler->final_learning_rate_frac = final_learning_rate_frac; -} - -void lr_scheduler_init_triangular(CyclicTriangularLearningRateScheduler *scheduler, float min_lr, float max_lr, int step_size) { - scheduler->min_lr = min_lr; - scheduler->max_lr = max_lr; - scheduler->step_size = step_size; -} - -void lr_scheduler_init_constant(ConstantLearningRateScheduler *scheduler, float learning_rate) { - scheduler->learning_rate = learning_rate; -} - #endif // SCHEDULERS_H \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index d311462..e9421a0 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1385,7 +1385,7 @@ int main(int argc, char *argv[]) { 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 if (argv[i][1] == 'k') { lr_scheduler_type = get_lr_scheduler_type_from_name(argv[i + 1]); } + else if (argv[i][1] == 'k') { lr_scheduler_type = get_lr_scheduler_type_from_name(argv[i+1]); } else { error_usage(); } } // should do a bit more error checking here @@ -1553,26 +1553,8 @@ int main(int argc, char *argv[]) { tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); // set up learning rate scheduler - CosineLearningRateScheduler lr_scheduler_cosine; - LinearLearningRateScheduler lr_scheduler_linear; - CyclicTriangularLearningRateScheduler lr_scheduler_triangular; - ConstantLearningRateScheduler lr_scheduler_constant; - if (lr_scheduler_type == LR_SCHEDULER_COSINE) { - lr_scheduler_init_cosine(&lr_scheduler_cosine, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); - } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { - lr_scheduler_init_linear(&lr_scheduler_linear, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); - } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { - // Hardcode some reasonable defaults for now - float min_lr = learning_rate / 10.0f; - float max_lr = learning_rate; - int step_size = train_num_batches / 4; - lr_scheduler_init_triangular(&lr_scheduler_triangular, min_lr, max_lr, step_size); - } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { - lr_scheduler_init_constant(&lr_scheduler_constant, learning_rate); - } else { - printf("Unknown learning rate scheduler type\n"); - exit(EXIT_FAILURE); - } + LearningRateScheduler lr_scheduler; + lr_scheduler_init(&lr_scheduler, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); // some memory for generating samples from the model int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); @@ -1733,19 +1715,7 @@ int main(int argc, char *argv[]) { // average the loss and the gradients between all processes gpt2_multi_gpu_loss_reduce(&model, &multi_gpu_config); // fetch the next learning rate - float step_learning_rate; - if (lr_scheduler_type == LR_SCHEDULER_COSINE) { - step_learning_rate = get_learning_rate_cosine(&lr_scheduler_cosine, step); - } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { - step_learning_rate = get_learning_rate_linear(&lr_scheduler_linear, step); - } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { - step_learning_rate = get_learning_rate_triangular(&lr_scheduler_triangular, step); - } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { - step_learning_rate = get_learning_rate_constant(&lr_scheduler_constant, step); - } else { - printf("Unknown learning rate scheduler type\n"); - exit(EXIT_FAILURE); - } + float step_learning_rate = get_learning_rate(lr_scheduler_type, &lr_scheduler, step); // update the model parameters float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, 1.0f, step+1, &multi_gpu_config); // zero out the gradients for the next iteration From 3e9fb6cb368fc9bb50f5b201dff7c12546dc9e19 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Wed, 19 Jun 2024 20:51:27 +0200 Subject: [PATCH 20/47] Fix compile issues - add comments --- llmc/schedulers.h | 42 ++++++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index b82dbdf..f24626b 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -7,6 +7,7 @@ Implements various learning rate schedulers. #include #include +#include typedef enum { LR_SCHEDULER_COSINE, @@ -62,24 +63,6 @@ void lr_scheduler_init(LearningRateScheduler *scheduler, float learning_rate, in // Learning rate scheduler functions // -// switch to the appropriate learning rate scheduler -float get_learning_rate(LRSchedulerType lr_scheduler_type, LearningRateScheduler *scheduler, int step) { - float step_learning_rate; - if (lr_scheduler_type == LR_SCHEDULER_COSINE) { - step_learning_rate = get_learning_rate_cosine(scheduler, step); - } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { - step_learning_rate = get_learning_rate_linear(scheduler, step); - } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { - step_learning_rate = get_learning_rate_triangular(scheduler, step); - } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { - step_learning_rate = get_learning_rate_constant(scheduler, step); - } else { - printf("Unknown learning rate scheduler type\n"); - exit(EXIT_FAILURE); - } - return step_learning_rate; -} - // cosine learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac float get_learning_rate_cosine(LearningRateScheduler *scheduler, int step) { float lr = scheduler->learning_rate; @@ -113,6 +96,7 @@ float get_learning_rate_linear(LearningRateScheduler *scheduler, int step) { // cyclic triangular learning rate schedule: linearly increase LR from min LR to max LR, then linearly decrease LR to min LR (repeat) // currently hardcoded to support only a single cycle float get_learning_rate_triangular(LearningRateScheduler *scheduler, int step) { + // warmup_iterations <- not used. int step_size = scheduler->train_num_batches / 2; // number of steps in half a cycle float min_lr = scheduler->learning_rate * scheduler->final_learning_rate_frac; float max_lr = scheduler->learning_rate; @@ -125,7 +109,29 @@ float get_learning_rate_triangular(LearningRateScheduler *scheduler, int step) { // constant learning rate schedule float get_learning_rate_constant(LearningRateScheduler *scheduler, int step) { + // warmup_iterations <- not used. + // train_num_batches <- not used. + // final_learning_rate_frac <- not used. return scheduler->learning_rate; } +// switch to the appropriate learning rate scheduler +float get_learning_rate(LRSchedulerType lr_scheduler_type, LearningRateScheduler *scheduler, int step) { + float step_learning_rate; + if (lr_scheduler_type == LR_SCHEDULER_COSINE) { + step_learning_rate = get_learning_rate_cosine(scheduler, step); + } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { + step_learning_rate = get_learning_rate_linear(scheduler, step); + } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { + step_learning_rate = get_learning_rate_triangular(scheduler, step); + } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { + step_learning_rate = get_learning_rate_constant(scheduler, step); + } else { + printf("Unknown learning rate scheduler type\n"); + exit(EXIT_FAILURE); + } + return step_learning_rate; +} + + #endif // SCHEDULERS_H \ No newline at end of file From 716808d9b93762f21db31e8704ac620997ea87f4 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Thu, 20 Jun 2024 16:09:28 +0200 Subject: [PATCH 21/47] Use C not C++ string header --- llmc/schedulers.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index f24626b..927edeb 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -7,7 +7,7 @@ Implements various learning rate schedulers. #include #include -#include +#include typedef enum { LR_SCHEDULER_COSINE, From 96ebb1629849544878dc3278bd0dc66fc5545a71 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 21 Jun 2024 00:29:47 +0000 Subject: [PATCH 22/47] simplify lr schedulers --- llmc/schedulers.h | 81 +++++++---------------------------------------- train_gpt2.cu | 13 ++++---- 2 files changed, 18 insertions(+), 76 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index 927edeb..44e9562 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -2,68 +2,29 @@ Implements various learning rate schedulers. */ #ifndef SCHEDULERS_H - #define SCHEDULERS_H #include #include #include -typedef enum { - LR_SCHEDULER_COSINE, - LR_SCHEDULER_LINEAR, - LR_SCHEDULER_TRIANGULAR, - LR_SCHEDULER_CONSTANT, - NUM_LR_SCHEDULERS // To keep track of the number of schedulers -} LRSchedulerType; - -const char* lr_scheduler_names[] = { - "cosine", - "linear", - "triangular", - "constant", -}; - -const char* get_lr_scheduler_name(LRSchedulerType type) { - if (type < 0 || type >= NUM_LR_SCHEDULERS) { - exit(EXIT_FAILURE); - } - return lr_scheduler_names[type]; -} - -LRSchedulerType get_lr_scheduler_type_from_name(const char* name) { - for (int i = 0; i < NUM_LR_SCHEDULERS; ++i) { - if (strcmp(name, lr_scheduler_names[i]) == 0) { - return (LRSchedulerType)i; - } - } - printf("Warning: Unknown learning rate scheduler name: %s\n. Using cosine as default.", name); - return LR_SCHEDULER_COSINE; // Default to cosine if not found -} - -// -// Learning rate scheduler structs and init -// - typedef struct { + const char* type; float learning_rate; int warmup_iterations; int train_num_batches; float final_learning_rate_frac; } LearningRateScheduler; -void lr_scheduler_init(LearningRateScheduler *scheduler, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { +void lr_scheduler_init(LearningRateScheduler *scheduler, const char* scheduler_type, float learning_rate, int warmup_iterations, int train_num_batches, float final_learning_rate_frac) { + scheduler->type = scheduler_type; scheduler->learning_rate = learning_rate; scheduler->warmup_iterations = warmup_iterations; scheduler->train_num_batches = train_num_batches; scheduler->final_learning_rate_frac = final_learning_rate_frac; } -// -// Learning rate scheduler functions -// - -// cosine learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac +// cosine: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac float get_learning_rate_cosine(LearningRateScheduler *scheduler, int step) { float lr = scheduler->learning_rate; if (step < scheduler->warmup_iterations) { @@ -79,7 +40,7 @@ float get_learning_rate_cosine(LearningRateScheduler *scheduler, int step) { return lr; } -// linear warmup learning rate schedule: warmup linearly to max LR, then decay linearly to LR * final_learning_rate_frac +// linear: warmup linearly to max LR, then decay linearly to LR * final_learning_rate_frac float get_learning_rate_linear(LearningRateScheduler *scheduler, int step) { float lr = scheduler->learning_rate; if (step < scheduler->warmup_iterations) { @@ -93,38 +54,19 @@ float get_learning_rate_linear(LearningRateScheduler *scheduler, int step) { return lr; } -// cyclic triangular learning rate schedule: linearly increase LR from min LR to max LR, then linearly decrease LR to min LR (repeat) -// currently hardcoded to support only a single cycle -float get_learning_rate_triangular(LearningRateScheduler *scheduler, int step) { - // warmup_iterations <- not used. - int step_size = scheduler->train_num_batches / 2; // number of steps in half a cycle - float min_lr = scheduler->learning_rate * scheduler->final_learning_rate_frac; - float max_lr = scheduler->learning_rate; - - int cycle_index = 1 + step / (2 * step_size); // tells us which cycle we are in, starting at 1 - float x = fabsf((float)step / step_size - 2 * cycle_index + 1); // goes from 0 to 1 to 0 - float lr = min_lr + (max_lr - min_lr) * fmaxf(0, (1 - x)); - return lr; -} - -// constant learning rate schedule +// constant float get_learning_rate_constant(LearningRateScheduler *scheduler, int step) { - // warmup_iterations <- not used. - // train_num_batches <- not used. - // final_learning_rate_frac <- not used. return scheduler->learning_rate; } -// switch to the appropriate learning rate scheduler -float get_learning_rate(LRSchedulerType lr_scheduler_type, LearningRateScheduler *scheduler, int step) { +// return the learning rate at a given step +float get_learning_rate(LearningRateScheduler *scheduler, int step) { float step_learning_rate; - if (lr_scheduler_type == LR_SCHEDULER_COSINE) { + if (strcmp(scheduler->type, "cosine") == 0) { step_learning_rate = get_learning_rate_cosine(scheduler, step); - } else if (lr_scheduler_type == LR_SCHEDULER_LINEAR) { + } else if (strcmp(scheduler->type, "linear") == 0) { step_learning_rate = get_learning_rate_linear(scheduler, step); - } else if (lr_scheduler_type == LR_SCHEDULER_TRIANGULAR) { - step_learning_rate = get_learning_rate_triangular(scheduler, step); - } else if (lr_scheduler_type == LR_SCHEDULER_CONSTANT) { + } else if (strcmp(scheduler->type, "constant") == 0) { step_learning_rate = get_learning_rate_constant(scheduler, step); } else { printf("Unknown learning rate scheduler type\n"); @@ -133,5 +75,4 @@ float get_learning_rate(LRSchedulerType lr_scheduler_type, LearningRateScheduler return step_learning_rate; } - #endif // SCHEDULERS_H \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index e9421a0..961dbc3 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1284,7 +1284,7 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename // ---------------------------------------------------------------------------- // CLI, poor man's argparse -// unclaimed flags lol: k,p +// unclaimed flags lol: p void error_usage() { fprintf(stderr, "Usage: ./train_gpt2cu [options]\n"); @@ -1334,7 +1334,7 @@ int main(int argc, char *argv[]) { const char* train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; const char* val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; const char* load_filename = "gpt2_124M_bf16.bin"; // bf16 weights of the model - LRSchedulerType lr_scheduler_type = LR_SCHEDULER_COSINE; + const char* lr_scheduler_type = "cosine"; const char* output_log_dir = NULL; int checkpoint_every = 0; // write optimization checkpoints every how many steps? int resume = 0; // resume the optimization, if one is found inside output_log_dir? @@ -1385,7 +1385,7 @@ int main(int argc, char *argv[]) { 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 if (argv[i][1] == 'k') { lr_scheduler_type = get_lr_scheduler_type_from_name(argv[i+1]); } + else if (argv[i][1] == 'k') { lr_scheduler_type = argv[i+1]; } else { error_usage(); } } // should do a bit more error checking here @@ -1419,7 +1419,7 @@ 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("| LR scheduler | %-50s |\n", get_lr_scheduler_name(lr_scheduler_type)); + printf0("| LR scheduler | %-50s |\n", lr_scheduler_type); 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); @@ -1554,7 +1554,8 @@ int main(int argc, char *argv[]) { // set up learning rate scheduler LearningRateScheduler lr_scheduler; - lr_scheduler_init(&lr_scheduler, learning_rate, warmup_iterations, train_num_batches, final_learning_rate_frac); + lr_scheduler_init(&lr_scheduler, lr_scheduler_type, learning_rate, + warmup_iterations, train_num_batches, final_learning_rate_frac); // some memory for generating samples from the model int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); @@ -1715,7 +1716,7 @@ int main(int argc, char *argv[]) { // average the loss and the gradients between all processes gpt2_multi_gpu_loss_reduce(&model, &multi_gpu_config); // fetch the next learning rate - float step_learning_rate = get_learning_rate(lr_scheduler_type, &lr_scheduler, step); + float step_learning_rate = get_learning_rate(&lr_scheduler, step); // update the model parameters float grad_norm = gpt2_update(&model, step_learning_rate, 0.9f, 0.95f, 1e-8f, weight_decay, 1.0f, step+1, &multi_gpu_config); // zero out the gradients for the next iteration From 4b5fecc8b2b5e77e208318c998de51ae27781d24 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 21 Jun 2024 00:49:01 +0000 Subject: [PATCH 23/47] small touchups and add the wsd schedule --- llmc/schedulers.h | 26 ++++++++++++++++++++++++-- train_gpt2.cu | 2 +- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/llmc/schedulers.h b/llmc/schedulers.h index 44e9562..9ddc570 100644 --- a/llmc/schedulers.h +++ b/llmc/schedulers.h @@ -59,6 +59,26 @@ float get_learning_rate_constant(LearningRateScheduler *scheduler, int step) { return scheduler->learning_rate; } +// wsd schedule: warmup linearly, keep constant, last 20% decay using 1 - sqrt decay to final_frac (should be 0.0) +// https://arxiv.org/abs/2405.18392 +float get_learning_rate_wsd(LearningRateScheduler *scheduler, int step) { + int decay_point = (int)(0.8f * scheduler->train_num_batches); + float max_lr = scheduler->learning_rate; + float lr = max_lr; + if (step < scheduler->warmup_iterations) { + float decay_ratio = ((float)(step + 1)) / scheduler->warmup_iterations; + lr = max_lr * decay_ratio; + } else if (step < decay_point) { + // noop, keep lr constant + } else { + float decay_ratio = ((float)(step - decay_point)) / (scheduler->train_num_batches - decay_point); + assert(0.0f <= decay_ratio && decay_ratio <= 1.0f); + float min_lr = max_lr * scheduler->final_learning_rate_frac; + return min_lr + (1.0f - sqrtf(decay_ratio)) * (max_lr - min_lr); + } + return lr; +} + // return the learning rate at a given step float get_learning_rate(LearningRateScheduler *scheduler, int step) { float step_learning_rate; @@ -68,11 +88,13 @@ float get_learning_rate(LearningRateScheduler *scheduler, int step) { step_learning_rate = get_learning_rate_linear(scheduler, step); } else if (strcmp(scheduler->type, "constant") == 0) { step_learning_rate = get_learning_rate_constant(scheduler, step); + } else if (strcmp(scheduler->type, "wsd") == 0) { + step_learning_rate = get_learning_rate_wsd(scheduler, step); } else { - printf("Unknown learning rate scheduler type\n"); + fprintf(stderr, "Unknown learning rate scheduler type: %s\n", scheduler->type); exit(EXIT_FAILURE); } return step_learning_rate; } -#endif // SCHEDULERS_H \ No newline at end of file +#endif // SCHEDULERS_H diff --git a/train_gpt2.cu b/train_gpt2.cu index 961dbc3..e07b431 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -21,7 +21,7 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. #include "llmc/dataloader.h" // defines: manual_seed, normal_ (same as torch.manual_seed and torch.normal) #include "llmc/rand.h" -// defines learning rate schedulers +// defines: lr_scheduler_init, get_learning_rate #include "llmc/schedulers.h" // defines: sample_softmax, random_f32 #include "llmc/sampler.h" From 61ed938dfb9c35272836e3c2fe4a6672a72c6b0a Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 21 Jun 2024 02:40:11 +0000 Subject: [PATCH 24/47] delete spurious check it's fine --- train_gpt2.cu | 6 ------ 1 file changed, 6 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index f9e846a..f4d7ca7 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1388,12 +1388,6 @@ int main(int argc, char *argv[]) { if (output_log_dir != NULL) { assert(strlen(output_log_dir) < 400); // careful bunch of hardcoded snprintf around this } - // check if output_log_dir does not exist or is a file - struct stat info; - if (output_log_dir != NULL && (stat(output_log_dir, &info ) != 0 || !(info.st_mode & S_IFDIR))) { - fprintf(stderr, "-o \"%s\" does not exist or is a file - are you specifying a file instead of dir?\n", output_log_dir); - exit(EXIT_FAILURE); - } int tokens_per_fwdbwd = B * T * multi_gpu_config.num_processes; // one micro-batch processes this many tokens // calculate sensible default for total batch size as assuming no gradient accumulation if (total_batch_size == -1) { total_batch_size = tokens_per_fwdbwd; } From 72a21588123db32b158f3006ff879d271c95d8fc Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 21 Jun 2024 15:08:30 +0000 Subject: [PATCH 25/47] remove a spurious print that is making the table a bit uglier --- train_gpt2.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 20eea11..a03bdbb 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -97,7 +97,6 @@ void set_zero_configs(MultiGpuConfig* multi_gpu_config, int zero_stage, size_t t multi_gpu_config->zero_stage = 0; } else { - printf0("| Zero Stage1 is enabled |\n"); multi_gpu_config->zero_stage = 1; multi_gpu_config->shard_num_parameters = total_parameters / multi_gpu_config->num_processes; } From 71f10e796070fe51ac049d0e7f285e48d4812063 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 22 Jun 2024 14:32:51 +0200 Subject: [PATCH 26/47] Rename gpt2 backward --- profile_gpt2.cu | 2 +- test_gpt2.cu | 6 +++--- train_gpt2.cu | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/profile_gpt2.cu b/profile_gpt2.cu index 8a7ddbf..9b30018 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -55,7 +55,7 @@ int main(int argc, char *argv[]) { // do a training step gpt2_forward(&model, x, y, B, T); gpt2_zero_grad(&model); - gpt2_backward(&model, x, true); + gpt2_backward_and_reduce(&model, x, true); gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, 1.f, 1, &multi_gpu_config); cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings diff --git a/test_gpt2.cu b/test_gpt2.cu index acf0a5d..491d287 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -212,7 +212,7 @@ int main(int argc, char *argv[]) { clock_gettime(CLOCK_MONOTONIC, &start); gpt2_forward(&model, x, y, B, T); gpt2_zero_grad(&model); - gpt2_backward(&model, x, true); + gpt2_backward_and_reduce(&model, x, true); clock_gettime(CLOCK_MONOTONIC, &end); double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9; @@ -328,7 +328,7 @@ int main(int argc, char *argv[]) { dataloader_next_batch(&loader); gpt2_forward(&model, loader.inputs, loader.targets, B, T); gpt2_zero_grad(&model); - gpt2_backward(&model, loader.inputs, true); + gpt2_backward_and_reduce(&model, loader.inputs, true); gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+11, &multi_gpu_config); losses[step] = model.mean_loss; tokens[step] = loader.inputs[0]; @@ -343,7 +343,7 @@ int main(int argc, char *argv[]) { dataloader_next_batch(&loader); gpt2_forward(&model, loader.inputs, loader.targets, B, T); gpt2_zero_grad(&model); - gpt2_backward(&model, loader.inputs, true); + gpt2_backward_and_reduce(&model, loader.inputs, true); gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+11, &multi_gpu_config); if(loader.inputs[0] != tokens[step]) { diff --git a/train_gpt2.cu b/train_gpt2.cu index 20eea11..db37876 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -722,7 +722,7 @@ void gpt2_zero_grad(GPT2 *model) { cudaCheck(cudaDeviceSynchronize()); } -void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { +void gpt2_backward_and_reduce(GPT2 *model, int* inputs, bool last_step) { NVTX_RANGE_FN(); // double check we forwarded previously, with targets if (model->mean_loss == -1.0f) { @@ -1702,7 +1702,7 @@ int main(int argc, char *argv[]) { gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T, grad_accum_steps); lossf += model.mean_loss; // the mean_loss was normalized by grad_accum_steps inside gpt2_forward // backward pass. all model params accumulate gradients with += inside this inner loop - gpt2_backward(&model, train_loader.inputs, micro_step == grad_accum_steps - 1); + gpt2_backward_and_reduce(&model, train_loader.inputs, micro_step == grad_accum_steps - 1); } // override the mean loss, accounting for the gradient accumulation loop // this is esp important to do here in multigpu update below, where model.mean_loss gets allreduced From 98e928f1d828383555c4746bd027be7df8c6b11c Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 23 Jun 2024 01:14:42 +0000 Subject: [PATCH 27/47] minor fixes for disk device io --- llmc/cuda_common.h | 3 +++ train_gpt2.cu | 5 +---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/llmc/cuda_common.h b/llmc/cuda_common.h index 2d973e4..5f031cb 100644 --- a/llmc/cuda_common.h +++ b/llmc/cuda_common.h @@ -118,6 +118,9 @@ class NvtxRange { }; #define NVTX_RANGE_FN() NvtxRange nvtx_range(__FUNCTION__) +// ---------------------------------------------------------------------------- +// Utilities to Read & Write between CUDA memory <-> files + // copy num_bytes from device pointer src into file dest, using double buffering running on the given stream. inline void device_to_file(FILE* dest, void* src, size_t num_bytes, size_t buffer_size, cudaStream_t stream) { // allocate pinned buffer for faster, async transfer diff --git a/train_gpt2.cu b/train_gpt2.cu index 4460e74..574c95c 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -387,7 +387,7 @@ void gpt2_write_to_checkpoint(GPT2 *model, const char* checkpoint_path) { fwriteCheck(model_header, sizeof(int), 256, model_file); // write the parameters device_to_file(model_file, model->params_memory, model->num_parameters_bytes, - 1024*1024*32, main_stream); + IO_BUF_SIZE, main_stream); // close file, we're done fcloseCheck(model_file); } @@ -1224,19 +1224,16 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float))); } - if(use_master_weights == 1 && !model->use_master_weights) { printf0("Warning: Master weights are present in state, but not enabled for current run."); } else if (use_master_weights == 0 && model->use_master_weights) { printf0("Error: Master weights requested, but not present in state file."); exit(EXIT_FAILURE); } - if (model->master_weights == NULL && use_master_weights == 1) { printf0("allocating %zu MiB for master copy of params\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->master_weights, shard_num_parameters * sizeof(float))); } - file_to_device(model->m_memory, state_file, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); file_to_device(model->v_memory, state_file, shard_num_parameters * sizeof(float), IO_BUF_SIZE, main_stream); if(model->use_master_weights) { From 0addb503ef009a6eac67a919e52e51038f103195 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 22 Jun 2024 17:16:35 +0200 Subject: [PATCH 28/47] Use NCCL where possible --- llmc/zero.cuh | 84 +++++++++++++++++++++++++++++++-------------------- train_gpt2.cu | 21 +++++++------ 2 files changed, 63 insertions(+), 42 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 160dae7..25c1e2d 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -50,11 +50,9 @@ void mpi_check(int status, const char *file, int line) { #endif // MULTI_GPU // ---------------------------------------------------------------------------- -// MPI / multi-processing setup - // Parameters specific to training on multiple GPUs. typedef struct { - int process_rank; // Rank of this process among all MPI processes. 0 if no multi-GPU. + int process_rank; // Rank of this process among all processes. 0 if no multi-GPU. int num_processes; // Total number of processes. 1 if no multi-GPU. int local_device_idx; // This process GPU index on current machine. 0 if no multi-GPU. @@ -69,6 +67,7 @@ typedef struct { ncclComm_t nccl_comm; // NCCL communication primitive, used for collective multi-GPU work. cudaStream_t nccl_stream; // CUDA Stream to perform NCCL operations. cudaEvent_t compute_nccl_sync; // Event used to synchronize NCCL with the compute + float* unified_buffer; #endif } MultiGpuConfig; @@ -77,39 +76,55 @@ typedef struct { // Processes on the same machines use different GPU indicies. Processes on other machines don't. // Copied from NCCL examples: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/examples.html#example-2-one-device-per-process-or-thread int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { - char hostname[1024]; - hostname[1023] = '\0'; - // All processes on the same machine will share the same hostname. - gethostname(hostname, 1023); - for (int i=0; i < 1024; i++) { - if (hostname[i] == '.') { - hostname[i] = '\0'; - break; + char hostname[1024]; + hostname[1023] = '\0'; + // All processes on the same machine will share the same hostname. + gethostname(hostname, 1023); + for (int i=0; i < 1024; i++) { + if (hostname[i] == '.') { + hostname[i] = '\0'; + break; + } } - } - uint64_t hostname_hash = 5381u; - for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5u) + hostname_hash) ^ hostname[c]; } + uint64_t hostname_hash = 5381u; + for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5u) + hostname_hash) ^ hostname[c]; } - // Distribute all hostname hashes to all processes. - uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); - all_hostsname_hashes[process_rank] = hostname_hash; - mpiCheck(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, all_hostsname_hashes, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD)); + // Distribute all hostname hashes to all processes. + // Allocate memory for hostname hashes on the GPU + uint64_t* all_hostsname_hashes; + cudaMalloc(&all_hostsname_hashes, num_processes * sizeof(uint64_t)); - // Identify which GPU we need to use. - int local_device_idx = 0; - for (int current_process = 0; current_process < num_processes; ++current_process) { - if (current_process == process_rank) { - // Found my gpu, local_device_idx now has my target GPU index. - break; - } - if (all_hostsname_hashes[current_process] == all_hostsname_hashes[process_rank]) { - // This process ID runs on the same machine, but it's not me, skip this GPU - local_device_idx++; - } - } + // Copy the local hostname hash to the appropriate location in the GPU memory + cudaMemcpy(&all_hostsname_hashes[process_rank], &hostname_hash, sizeof(uint64_t), cudaMemcpyHostToDevice); - free(all_hostsname_hashes); - return local_device_idx; + // Perform the all-gather operation with NCCL + ncclAllGather(&hostname_hash, all_hostsname_hashes, sizeof(uint64_t), ncclUint64, nccl_comm, nccl_stream); + + // Synchronize the stream to ensure the all-gather operation is complete + cudaStreamSynchronize(nccl_stream); + + // Allocate CPU memory to copy the gathered results back + uint64_t* all_hostsname_hashes_cpu = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); + + // Copy the gathered results back to the CPU + cudaMemcpy(all_hostsname_hashes_cpu, all_hostsname_hashes, num_processes * sizeof(uint64_t), cudaMemcpyDeviceToHost); + + // Identify which GPU we need to use. + int local_device_idx = 0; + for (int current_process = 0; current_process < num_processes; ++current_process) { + if (current_process == process_rank) { + // Found my gpu, local_device_idx now has my target GPU index. + break; + } + if (all_hostsname_hashes[current_process] == all_hostsname_hashes[process_rank]) { + // This process ID runs on the same machine, but it's not me, skip this GPU + local_device_idx++; + } + } + + cudaFree(all_hostsname_hashes); + free(all_hostsname_hashes_cpu); + return local_device_idx; } #endif @@ -133,6 +148,7 @@ MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) { cudaCheck(cudaEventCreate(&result.compute_nccl_sync, cudaEventDisableTiming)); nvtxNameCudaStreamA(result.nccl_stream, "nccl stream"); nvtxNameCudaEventA(result.compute_nccl_sync, "nccl compute sync"); + cudaCheck(cudaMallocManaged(&result.unified_buffer, sizeof(float))); return result; #else printf("Multi-GPU support is disabled. Using a single GPU.\n"); @@ -157,8 +173,10 @@ void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) { void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config) { #ifdef MULTI_GPU if (multi_gpu_config->num_processes > 1) { - mpiCheck(MPI_Barrier(MPI_COMM_WORLD)); + float* unified_buffer = multi_gpu_config->unified_buffer; + ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, main_stream)); } + cudaCheck(cudaDeviceSynchronize()); #endif } diff --git a/train_gpt2.cu b/train_gpt2.cu index 574c95c..230100a 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -893,12 +893,15 @@ void gpt2_backward(GPT2 *model, int* inputs, bool last_step) { } // Compute sum of a single CPU value across all GPU processes. No-op when multi-GPU is disabled. -float multi_gpu_cpu_float_sum(float value) { +float multi_gpu_cpu_float_sum(float value, MultiGpuConfig* multi_gpu_config) { #ifdef MULTI_GPU - // note MPI doesn't support all reduce with mean, only sum - float result; - mpiCheck(MPI_Allreduce(&value, &result, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD)); - return result; + if (multi_gpu_config->num_processes == 1) return value; + + float* unified_buffer = multi_gpu_config->unified_buffer; + *unified_buffer = value; + ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, main_stream)); + cudaCheck(cudaDeviceSynchronize()); + return *unified_buffer; #else return value; #endif @@ -912,7 +915,7 @@ void gpt2_multi_gpu_loss_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) { // If there's only one process, there is nothing to do if (multi_gpu_config->num_processes == 1) { return; } // Average all losses. - model->accumulated_mean_loss = multi_gpu_cpu_float_sum(model->mean_loss) / multi_gpu_config->num_processes; + model->accumulated_mean_loss = multi_gpu_cpu_float_sum(model->mean_loss, multi_gpu_config) / multi_gpu_config->num_processes; #endif cudaCheck(cudaDeviceSynchronize()); } @@ -990,7 +993,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream); cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); // further sum the (partial) squared norm across all GPUs (see comment ^1 above) - grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu); + grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu, multi_gpu_config); } else { // in regular DDP, backward has averaged the gradients across all GPUs // so each GPU can compute the squared norm over the whole grad vector, with no added comms needed @@ -1579,7 +1582,7 @@ int main(int argc, char *argv[]) { val_loss += model.mean_loss; } val_loss /= val_num_batches; - val_loss = multi_gpu_cpu_float_sum(val_loss) / multi_gpu_config.num_processes; + val_loss = multi_gpu_cpu_float_sum(val_loss, &multi_gpu_config) / multi_gpu_config.num_processes; printf0("val loss %f\n", val_loss); logger_log_val(&logger, step, val_loss); } @@ -1598,7 +1601,7 @@ int main(int argc, char *argv[]) { eval_acc_norm += (float)correct; } // 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); + eval_acc_norm = multi_gpu_cpu_float_sum(eval_acc_norm, &multi_gpu_config); 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 / eval_loader.num_examples); } From 52db158367af337d3c3ad22252248c069a0e3281 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 22 Jun 2024 20:22:06 +0000 Subject: [PATCH 29/47] Can not call set device after NCCL init --- llmc/zero.cuh | 26 +++++--------------------- 1 file changed, 5 insertions(+), 21 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 25c1e2d..084b695 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -90,24 +90,9 @@ int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5u) + hostname_hash) ^ hostname[c]; } // Distribute all hostname hashes to all processes. - // Allocate memory for hostname hashes on the GPU - uint64_t* all_hostsname_hashes; - cudaMalloc(&all_hostsname_hashes, num_processes * sizeof(uint64_t)); - - // Copy the local hostname hash to the appropriate location in the GPU memory - cudaMemcpy(&all_hostsname_hashes[process_rank], &hostname_hash, sizeof(uint64_t), cudaMemcpyHostToDevice); - - // Perform the all-gather operation with NCCL - ncclAllGather(&hostname_hash, all_hostsname_hashes, sizeof(uint64_t), ncclUint64, nccl_comm, nccl_stream); - - // Synchronize the stream to ensure the all-gather operation is complete - cudaStreamSynchronize(nccl_stream); - - // Allocate CPU memory to copy the gathered results back - uint64_t* all_hostsname_hashes_cpu = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); - - // Copy the gathered results back to the CPU - cudaMemcpy(all_hostsname_hashes_cpu, all_hostsname_hashes, num_processes * sizeof(uint64_t), cudaMemcpyDeviceToHost); + uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); + all_hostsname_hashes[process_rank] = hostname_hash; + mpiCheck(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, all_hostsname_hashes, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD)); // Identify which GPU we need to use. int local_device_idx = 0; @@ -122,8 +107,7 @@ int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { } } - cudaFree(all_hostsname_hashes); - free(all_hostsname_hashes_cpu); + free(all_hostsname_hashes); return local_device_idx; } #endif @@ -174,7 +158,7 @@ void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config) { #ifdef MULTI_GPU if (multi_gpu_config->num_processes > 1) { float* unified_buffer = multi_gpu_config->unified_buffer; - ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, main_stream)); + ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); } cudaCheck(cudaDeviceSynchronize()); #endif From ac72875c24dcff74efc27ae7d723fbe195ed1e4b Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 23 Jun 2024 10:39:55 +0000 Subject: [PATCH 30/47] Use NCCL stream --- train_gpt2.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 230100a..be2aedd 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -899,7 +899,7 @@ float multi_gpu_cpu_float_sum(float value, MultiGpuConfig* multi_gpu_config) { float* unified_buffer = multi_gpu_config->unified_buffer; *unified_buffer = value; - ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, main_stream)); + ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); cudaCheck(cudaDeviceSynchronize()); return *unified_buffer; #else From 1bd737546c20011069c02f38f4176da6ffd8a16a Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 23 Jun 2024 10:42:19 +0000 Subject: [PATCH 31/47] Mini refactor --- llmc/zero.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 084b695..e64c7e6 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -157,8 +157,7 @@ void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) { void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config) { #ifdef MULTI_GPU if (multi_gpu_config->num_processes > 1) { - float* unified_buffer = multi_gpu_config->unified_buffer; - ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); + ncclCheck(ncclAllReduce(multi_gpu_config->unified_buffer, multi_gpu_config->unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream)); } cudaCheck(cudaDeviceSynchronize()); #endif From 08c055cd72db49eac39b40285e84dc289eaf3728 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 23 Jun 2024 10:46:54 +0000 Subject: [PATCH 32/47] Free up unified buffer mem --- llmc/zero.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index e64c7e6..7bcf08c 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -150,6 +150,7 @@ void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) { ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm)); cudaCheck(cudaStreamDestroy(multi_gpu_config->nccl_stream)); cudaCheck(cudaEventDestroy(multi_gpu_config->compute_nccl_sync)); + cudaCheck(cudaFree(multi_gpu_config->unified_buffer)); mpiCheck(MPI_Finalize()); #endif } From 5500b61da59f1c476640f636c348df657931ead4 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 23 Jun 2024 16:29:54 +0000 Subject: [PATCH 33/47] fs working --- llmc/zero.cuh | 94 ++++++++++++++++++--------------------------------- train_gpt2.cu | 11 ++++-- 2 files changed, 41 insertions(+), 64 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 7bcf08c..86996ef 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -12,7 +12,6 @@ Utilities for ZeRO sharding #include #ifdef MULTI_GPU -#include #include #endif @@ -36,17 +35,6 @@ void nccl_check(ncclResult_t status, const char *file, int line) { } #define ncclCheck(err) (nccl_check(err, __FILE__, __LINE__)) -void mpi_check(int status, const char *file, int line) { - if (status != MPI_SUCCESS) { - char mpi_error[4096]; - int mpi_error_len = 0; - assert(MPI_Error_string(status, &mpi_error[0], &mpi_error_len) == MPI_SUCCESS); - printf("[MPI ERROR] at file %s:%d:\n%.*s\n", file, line, mpi_error_len, mpi_error); - exit(EXIT_FAILURE); - } -} -#define mpiCheck(err) (mpi_check(err, __FILE__, __LINE__)) - #endif // MULTI_GPU // ---------------------------------------------------------------------------- @@ -71,61 +59,46 @@ typedef struct { #endif } MultiGpuConfig; +MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank) { #ifdef MULTI_GPU -// Determine which GPU this process should use. -// Processes on the same machines use different GPU indicies. Processes on other machines don't. -// Copied from NCCL examples: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/examples.html#example-2-one-device-per-process-or-thread -int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { - char hostname[1024]; - hostname[1023] = '\0'; - // All processes on the same machine will share the same hostname. - gethostname(hostname, 1023); - for (int i=0; i < 1024; i++) { - if (hostname[i] == '.') { - hostname[i] = '\0'; - break; - } - } - uint64_t hostname_hash = 5381u; - for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5u) + hostname_hash) ^ hostname[c]; } - - // Distribute all hostname hashes to all processes. - uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); - all_hostsname_hashes[process_rank] = hostname_hash; - mpiCheck(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, all_hostsname_hashes, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD)); - - // Identify which GPU we need to use. - int local_device_idx = 0; - for (int current_process = 0; current_process < num_processes; ++current_process) { - if (current_process == process_rank) { - // Found my gpu, local_device_idx now has my target GPU index. - break; - } - if (all_hostsname_hashes[current_process] == all_hostsname_hashes[process_rank]) { - // This process ID runs on the same machine, but it's not me, skip this GPU - local_device_idx++; - } - } - - free(all_hostsname_hashes); - return local_device_idx; -} -#endif - -MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) { -#ifdef MULTI_GPU - // Initialize MPI. MultiGpuConfig result; - mpiCheck(MPI_Init(argc, argv)); - mpiCheck(MPI_Comm_rank(MPI_COMM_WORLD, &result.process_rank)); - mpiCheck(MPI_Comm_size(MPI_COMM_WORLD, &result.num_processes)); - result.local_device_idx = multi_gpu_get_local_device_idx(result.process_rank, result.num_processes); + result.process_rank = process_rank; + result.num_processes = num_processes; + result.local_device_idx = process_rank % 8; cudaCheck(cudaSetDevice(result.local_device_idx)); ncclUniqueId nccl_id; + + FILE* idFile; + char dfs_path[256] = "/ephemeral/data/tokenizers"; + static char filename[256]; + snprintf(filename, sizeof(filename), "%s/ncclUniqueId.dat", dfs_path); + if (result.process_rank == 0) { ncclCheck(ncclGetUniqueId(&nccl_id)); + idFile = fopen(filename, "wb"); + assert(idFile != NULL); + fwrite(&nccl_id, sizeof(nccl_id), 1, idFile); + fclose(idFile); + // Construct the scp command + char command[1024]; + snprintf(command, sizeof(command), "scp %s ubuntu@h100-node-1-1:%s", filename, filename); + printf("Executing command: %s\n", command); + // Execute the scp command + int ret = system(command); + if (ret != 0) { + fprintf(stderr, "scp command failed with error code %d\n", ret); + exit(EXIT_FAILURE); + } + } else { // Other ranks wait until the file is available and read the unique ID + do { + printf("%d: Waiting for the file to be available\n", result.process_rank); + usleep(1000000); + idFile = fopen(filename, "rb"); + if (idFile != NULL) break; + } while (idFile == NULL); + freadCheck(&nccl_id, sizeof(nccl_id), 1, idFile); + fclose(idFile); } - mpiCheck(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD)); ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank)); cudaCheck(cudaStreamCreate(&result.nccl_stream)); // event without timing for maximum performance @@ -151,7 +124,6 @@ void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) { cudaCheck(cudaStreamDestroy(multi_gpu_config->nccl_stream)); cudaCheck(cudaEventDestroy(multi_gpu_config->compute_nccl_sync)); cudaCheck(cudaFree(multi_gpu_config->unified_buffer)); - mpiCheck(MPI_Finalize()); #endif } diff --git a/train_gpt2.cu b/train_gpt2.cu index be2aedd..b593b8d 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1324,8 +1324,6 @@ void error_usage() { // ---------------------------------------------------------------------------- // main training loop int main(int argc, char *argv[]) { - multi_gpu_config = multi_gpu_config_init(&argc, &argv); - // read in the (optional) command line arguments const char* train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; const char* val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; @@ -1352,10 +1350,12 @@ 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 int hellaswag_eval = 0; + int num_processes = 1; + int process_rank = 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 - if (strlen(argv[i]) != 2) { error_usage(); } // must be -x (one dash, one letter) + if (!(strlen(argv[i]) == 2 || strlen(argv[i]) == 3)) { error_usage(); } // must be -x (one dash, one or two letters) // read in the args if (argv[i][1] == 'i') { train_data_pattern = argv[i+1]; } else if (argv[i][1] == 'j') { val_data_pattern = argv[i+1]; } @@ -1382,8 +1382,13 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 'r') { recompute = atoi(argv[i+1]); } else if (argv[i][1] == 'h') { hellaswag_eval = atoi(argv[i+1]); } else if (argv[i][1] == 'k') { lr_scheduler_type = argv[i+1]; } + else if (argv[i][1] == 'p' && argv[i][2] == 'n') { num_processes = atoi(argv[i+1]); } + else if (argv[i][1] == 'p' && argv[i][2] == 'r') { process_rank = atoi(argv[i+1]); } else { error_usage(); } } + printf("num_processes = %d, process_rank = %d\n", num_processes, process_rank); + multi_gpu_config = multi_gpu_config_init(num_processes, process_rank); + // should do a bit more error checking here assert(warmup_iterations >= 0); if (output_log_dir != NULL) { From 1c6c5ff0b75c2de4ce0801d2650c72273bd9a3e9 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 23 Jun 2024 20:30:09 +0000 Subject: [PATCH 34/47] TCP working --- llmc/zero.cuh | 131 ++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 105 insertions(+), 26 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 86996ef..c19ff7b 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -5,6 +5,7 @@ Utilities for ZeRO sharding #ifndef LLMC_ZERO_CUH #define LLMC_ZERO_CUH +#include #include #include #include @@ -59,6 +60,18 @@ typedef struct { #endif } MultiGpuConfig; +#ifdef MULTI_GPU +void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int num_clients) { + for (int i = 0; i < num_clients; ++i) { + if (send(client_sockets[i], nccl_id, sizeof(*nccl_id), 0) == -1) { + printf("Failed to send nccl_id"); + exit(EXIT_FAILURE); + } + close(client_sockets[i]); + } +} +#endif + MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank) { #ifdef MULTI_GPU MultiGpuConfig result; @@ -68,36 +81,102 @@ MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank) { cudaCheck(cudaSetDevice(result.local_device_idx)); ncclUniqueId nccl_id; - FILE* idFile; - char dfs_path[256] = "/ephemeral/data/tokenizers"; - static char filename[256]; - snprintf(filename, sizeof(filename), "%s/ncclUniqueId.dat", dfs_path); - + int SERVER_PORT = 12345; + const char* SERVER_IP = "10.0.1.220"; if (result.process_rank == 0) { ncclCheck(ncclGetUniqueId(&nccl_id)); - idFile = fopen(filename, "wb"); - assert(idFile != NULL); - fwrite(&nccl_id, sizeof(nccl_id), 1, idFile); - fclose(idFile); - // Construct the scp command - char command[1024]; - snprintf(command, sizeof(command), "scp %s ubuntu@h100-node-1-1:%s", filename, filename); - printf("Executing command: %s\n", command); - // Execute the scp command - int ret = system(command); - if (ret != 0) { - fprintf(stderr, "scp command failed with error code %d\n", ret); + + int MAX_CLIENTS = num_processes - 1; + int client_sockets[MAX_CLIENTS]; + int num_clients = 0; + int server_socket, new_socket; + struct sockaddr_in address; + int addrlen = sizeof(address); + int opt = 1; + + // Create a TCP socket + if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { + printf("Socket failed"); exit(EXIT_FAILURE); } - } else { // Other ranks wait until the file is available and read the unique ID - do { - printf("%d: Waiting for the file to be available\n", result.process_rank); - usleep(1000000); - idFile = fopen(filename, "rb"); - if (idFile != NULL) break; - } while (idFile == NULL); - freadCheck(&nccl_id, sizeof(nccl_id), 1, idFile); - fclose(idFile); + + // set socket options + // SOL_SOCKET - means that option is configured at socket level + // SO_REUSEADDR - allows to bind to an address which is in a TIME_WAIT state (already used by another socket) - useful when restarting the server + // SO_REUSEPORT - allows to bind to the same port multiple times + if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)) < 0) { + printf("Setsockopt failed"); + exit(EXIT_FAILURE); + } + + address.sin_family = AF_INET; // IPv4 + address.sin_addr.s_addr = inet_addr(SERVER_IP); // alternatively use INADDR_ANY to bind to all interfaces, currently we only allow ethernet + address.sin_port = htons(SERVER_PORT); + + // Bind the socket to the address and port + if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) < 0) { + printf("Bind failed"); + exit(EXIT_FAILURE); + } + + // MAX_CLIENTS specifies the maximum number of clients that can be queued for this server + if (listen(server_socket, MAX_CLIENTS) < 0) { + printf("Listen failed"); + exit(EXIT_FAILURE); + } + + printf("Waiting for clients to connect...\n"); + while (num_clients < MAX_CLIENTS) { + if ((new_socket = accept(server_socket, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) { + printf("Accept failed"); + exit(EXIT_FAILURE); + } + client_sockets[num_clients++] = new_socket; + printf("Client %d connected\n", num_clients); + } + + send_nccl_id_to_clients(&nccl_id, client_sockets, num_clients); + printf("NCCL ID sent to all clients\n"); + + close(server_socket); + } else { + int num_attempts = 5; + int time_to_sleep = 2; + + int client_socket; + struct sockaddr_in serv_addr; + + // Create a TCP socket + if ((client_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { + printf("Socket creation error"); + exit(EXIT_FAILURE); + } + + // Set the server address and port + serv_addr.sin_family = AF_INET; + serv_addr.sin_port = htons(SERVER_PORT); + if (inet_pton(AF_INET, SERVER_IP, &serv_addr.sin_addr) <= 0) { + printf("Invalid address or address not supported"); + exit(EXIT_FAILURE); + } + + // Try to connect to the server - retry if connection fails + while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) { + printf("%d Connection failed, retrying in %d seconds\n", process_rank, time_to_sleep); + if (--num_attempts == 0) { + printf("Failed to connect to the server\n"); + exit(EXIT_FAILURE); + } + sleep(time_to_sleep); + } + + if (recv(client_socket, &nccl_id, sizeof(nccl_id), 0) <= 0) { + printf("Failed to receive nccl_id"); + exit(EXIT_FAILURE); + } + + printf("Received NCCL ID\n"); + close(client_socket); } ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank)); cudaCheck(cudaStreamCreate(&result.nccl_stream)); From 1bc3a5ef6e45bc821b3a22234cb44afb43e9d3f5 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sun, 23 Jun 2024 21:42:05 +0000 Subject: [PATCH 35/47] Add all 3 init methods --- llmc/zero.cuh | 148 +++++++++++++++++++++++++++++++++++++++++++++----- train_gpt2.cu | 16 ++++-- 2 files changed, 144 insertions(+), 20 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index c19ff7b..dc7ba69 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -16,6 +16,10 @@ Utilities for ZeRO sharding #include #endif +#ifdef USE_MPI +#include +#endif + // ---------------------------------------------------------------------------- // Multi-GPU related #ifdef MULTI_GPU @@ -36,6 +40,19 @@ void nccl_check(ncclResult_t status, const char *file, int line) { } #define ncclCheck(err) (nccl_check(err, __FILE__, __LINE__)) +#ifdef USE_MPI +void mpi_check(int status, const char *file, int line) { + if (status != MPI_SUCCESS) { + char mpi_error[4096]; + int mpi_error_len = 0; + assert(MPI_Error_string(status, &mpi_error[0], &mpi_error_len) == MPI_SUCCESS); + printf("[MPI ERROR] at file %s:%d:\n%.*s\n", file, line, mpi_error_len, mpi_error); + exit(EXIT_FAILURE); + } +} +#define mpiCheck(err) (mpi_check(err, __FILE__, __LINE__)) +#endif + #endif // MULTI_GPU // ---------------------------------------------------------------------------- @@ -70,23 +87,15 @@ void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int nu close(client_sockets[i]); } } -#endif -MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank) { -#ifdef MULTI_GPU - MultiGpuConfig result; - result.process_rank = process_rank; - result.num_processes = num_processes; - result.local_device_idx = process_rank % 8; - cudaCheck(cudaSetDevice(result.local_device_idx)); +ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) { ncclUniqueId nccl_id; - + // hardcode an arbitrary port number between 1024 and 49151 (registered ports) int SERVER_PORT = 12345; - const char* SERVER_IP = "10.0.1.220"; - if (result.process_rank == 0) { + if (result->process_rank == 0) { ncclCheck(ncclGetUniqueId(&nccl_id)); - int MAX_CLIENTS = num_processes - 1; + int MAX_CLIENTS = result->num_processes - 1; int client_sockets[MAX_CLIENTS]; int num_clients = 0; int server_socket, new_socket; @@ -110,7 +119,7 @@ MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank) { } address.sin_family = AF_INET; // IPv4 - address.sin_addr.s_addr = inet_addr(SERVER_IP); // alternatively use INADDR_ANY to bind to all interfaces, currently we only allow ethernet + address.sin_addr.s_addr = inet_addr(server_ip); // alternatively use INADDR_ANY to bind to all interfaces, currently we only allow ethernet address.sin_port = htons(SERVER_PORT); // Bind the socket to the address and port @@ -155,14 +164,14 @@ MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank) { // Set the server address and port serv_addr.sin_family = AF_INET; serv_addr.sin_port = htons(SERVER_PORT); - if (inet_pton(AF_INET, SERVER_IP, &serv_addr.sin_addr) <= 0) { + if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) { printf("Invalid address or address not supported"); exit(EXIT_FAILURE); } // Try to connect to the server - retry if connection fails while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) { - printf("%d Connection failed, retrying in %d seconds\n", process_rank, time_to_sleep); + printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep); if (--num_attempts == 0) { printf("Failed to connect to the server\n"); exit(EXIT_FAILURE); @@ -178,6 +187,112 @@ MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank) { printf("Received NCCL ID\n"); close(client_socket); } + + return nccl_id; +} + +ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) { + // Works assuming that the filesystem is shared among all processes + ncclUniqueId nccl_id; + FILE* idFile; + static char filename[1024]; + snprintf(filename, sizeof(filename), "%s/ncclUniqueId.sync", fs_path); + + if (result->process_rank == 0) { + ncclCheck(ncclGetUniqueId(&nccl_id)); + idFile = fopen(filename, "wb"); + assert(idFile != NULL); + fwrite(&nccl_id, sizeof(nccl_id), 1, idFile); + fclose(idFile); + } else { + // Other ranks wait until the file is available and read the unique ID + do { + usleep(1000000); // 1 second + idFile = fopen(filename, "rb"); + if (idFile != NULL) break; + } while (idFile == NULL); + freadCheck(&nccl_id, sizeof(nccl_id), 1, idFile); + fclose(idFile); + } + + return nccl_id; +} + +#ifdef USE_MPI +// Determine which GPU this process should use. +// Processes on the same machines use different GPU indicies. Processes on other machines don't. +// Copied from NCCL examples: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/examples.html#example-2-one-device-per-process-or-thread +int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { + char hostname[1024]; + hostname[1023] = '\0'; + // All processes on the same machine will share the same hostname. + gethostname(hostname, 1023); + for (int i=0; i < 1024; i++) { + if (hostname[i] == '.') { + hostname[i] = '\0'; + break; + } + } + uint64_t hostname_hash = 5381u; + for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5u) + hostname_hash) ^ hostname[c]; } + + // Distribute all hostname hashes to all processes. + uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); + all_hostsname_hashes[process_rank] = hostname_hash; + mpiCheck(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, all_hostsname_hashes, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD)); + + // Identify which GPU we need to use. + int local_device_idx = 0; + for (int current_process = 0; current_process < num_processes; ++current_process) { + if (current_process == process_rank) { + // Found my gpu, local_device_idx now has my target GPU index. + break; + } + if (all_hostsname_hashes[current_process] == all_hostsname_hashes[process_rank]) { + // This process ID runs on the same machine, but it's not me, skip this GPU + local_device_idx++; + } + } + + free(all_hostsname_hashes); + return local_device_idx; +} +#endif + +#endif + +MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank, char* server_ip, char* fs_path, char* init_method) { +#ifdef MULTI_GPU + MultiGpuConfig result; + ncclUniqueId nccl_id; + if (strcmp(init_method, "mpi") != 0) { + result.process_rank = process_rank; + result.num_processes = num_processes; + result.local_device_idx = process_rank % 8; + if (strcmp(init_method, "tcp") == 0) { + nccl_id = get_nccl_id_via_tcp(&result, server_ip); + } else if (strcmp(init_method, "fs") == 0) { + nccl_id = get_nccl_id_via_fs(&result, fs_path); + } else { + printf("Invalid init method\n"); + exit(EXIT_FAILURE); + } + } else { + #ifdef USE_MPI + mpiCheck(MPI_Init(NULL, NULL)); + mpiCheck(MPI_Comm_rank(MPI_COMM_WORLD, &result.process_rank)); + mpiCheck(MPI_Comm_size(MPI_COMM_WORLD, &result.num_processes)); + result.local_device_idx = multi_gpu_get_local_device_idx(result.process_rank, result.num_processes); + if (result.process_rank == 0) { + ncclCheck(ncclGetUniqueId(&nccl_id)); + } + mpiCheck(MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD)); + #else + printf("MPI support is disabled. Please enable MPI support to use MPI init method.\n"); + exit(EXIT_FAILURE); + #endif + } + cudaCheck(cudaSetDevice(result.local_device_idx)); ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank)); cudaCheck(cudaStreamCreate(&result.nccl_stream)); // event without timing for maximum performance @@ -203,6 +318,9 @@ void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) { cudaCheck(cudaStreamDestroy(multi_gpu_config->nccl_stream)); cudaCheck(cudaEventDestroy(multi_gpu_config->compute_nccl_sync)); cudaCheck(cudaFree(multi_gpu_config->unified_buffer)); + #ifdef USE_MPI + mpiCheck(MPI_Finalize()); + #endif #endif } diff --git a/train_gpt2.cu b/train_gpt2.cu index b593b8d..215caca 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1350,12 +1350,16 @@ 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 int hellaswag_eval = 0; - int num_processes = 1; - int process_rank = 0; + // multi-node settings + int num_processes = 1; // this should be set by the slurm environment + int process_rank = 0; // this should be set by the slurm environment + char nccl_init_method[256] = "tcp"; // "tcp" or "fs" or "mpi" + char server_ip[256] = "-1"; // used if init_method set to "tcp" -> set to your server ip address + char fs_path[256] = "/tmp"; // used if init_method set to "fs" -> set to a shared filesystem path 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 - if (!(strlen(argv[i]) == 2 || strlen(argv[i]) == 3)) { error_usage(); } // must be -x (one dash, one or two letters) + if (!(strlen(argv[i]) == 2 || strlen(argv[i]) == 3)) { error_usage(); } // must be -x[y] (one dash, one or two letters) // read in the args if (argv[i][1] == 'i') { train_data_pattern = argv[i+1]; } else if (argv[i][1] == 'j') { val_data_pattern = argv[i+1]; } @@ -1382,12 +1386,14 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 'r') { recompute = atoi(argv[i+1]); } else if (argv[i][1] == 'h') { hellaswag_eval = atoi(argv[i+1]); } else if (argv[i][1] == 'k') { lr_scheduler_type = argv[i+1]; } + else if (argv[i][1] == 'p' && argv[i][2] == 'i') { strcpy(nccl_init_method, argv[i+1]); } + else if (argv[i][1] == 'p' && argv[i][2] == 'f') { strcpy(fs_path, argv[i+1]); } + else if (argv[i][1] == 'p' && argv[i][2] == 's') { strcpy(server_ip, argv[i+1]); } else if (argv[i][1] == 'p' && argv[i][2] == 'n') { num_processes = atoi(argv[i+1]); } else if (argv[i][1] == 'p' && argv[i][2] == 'r') { process_rank = atoi(argv[i+1]); } else { error_usage(); } } - printf("num_processes = %d, process_rank = %d\n", num_processes, process_rank); - multi_gpu_config = multi_gpu_config_init(num_processes, process_rank); + multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, server_ip, fs_path, nccl_init_method); // should do a bit more error checking here assert(warmup_iterations >= 0); From bfdea887ebdbbfe0c198bd097666a8e40b0391de Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 11:28:53 +0000 Subject: [PATCH 36/47] Add launch scripts - tested --- Makefile | 28 ++++--- llmc/zero.cuh | 8 +- scripts/multi_node/run_gpt2_124M_fs.sbatch | 84 ++++++++++++++++++++ scripts/multi_node/run_gpt2_124M_mpi.sh | 49 ++++++++++++ scripts/multi_node/run_gpt2_124M_tcp.sbatch | 85 +++++++++++++++++++++ 5 files changed, 243 insertions(+), 11 deletions(-) create mode 100755 scripts/multi_node/run_gpt2_124M_fs.sbatch create mode 100755 scripts/multi_node/run_gpt2_124M_mpi.sh create mode 100755 scripts/multi_node/run_gpt2_124M_tcp.sbatch diff --git a/Makefile b/Makefile index 9ea6866..200b81c 100644 --- a/Makefile +++ b/Makefile @@ -188,27 +188,35 @@ else endif endif -# Check if OpenMPI and NCCL are available, include them if so, for multi-GPU training +# Check if NCCL is available, include if so, for multi-GPU training ifeq ($(NO_MULTI_GPU), 1) - $(info → Multi-GPU (OpenMPI + NCCL) is manually disabled) + $(info → Multi-GPU (NCCL) is manually disabled) else ifneq ($(OS), Windows_NT) # Detect if running on macOS or Linux ifeq ($(SHELL_UNAME), Darwin) - $(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping OpenMPI + NCCL support) - else ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists) - $(info ✓ OpenMPI found, OK to train with multiple GPUs) - NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include - NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/ - NVCC_LDLIBS += -lmpi -lnccl + $(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping NCCL support) + else ifeq ($(shell dpkg -l | grep -q nccl && echo "exists"), exists) + $(info ✓ NCCL found, OK to train with multiple GPUs) + NVCC_LDLIBS += -lnccl NVCC_FLAGS += -DMULTI_GPU else - $(info ✗ OpenMPI is not found, disabling multi-GPU support) - $(info ---> On Linux you can try install OpenMPI with `sudo apt install openmpi-bin openmpi-doc libopenmpi-dev`) + $(info ✗ NCCL is not found, disabling multi-GPU support) + $(info ---> On Linux you can try install NCCL with `sudo apt install libnccl2 libnccl-dev`) endif endif endif +ifeq ($(USE_MPI), 1) + $(info → MPI is manually enabled) + NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include + NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/ + NVCC_FLAGS += -DUSE_MPI + NVCC_LDLIBS += -lmpi +else + $(info → MPI is manually disabled) +endif + # Precision settings, default to bf16 but ability to override PRECISION ?= BF16 VALID_PRECISIONS := FP32 FP16 BF16 diff --git a/llmc/zero.cuh b/llmc/zero.cuh index dc7ba69..e85e6b8 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -198,6 +198,12 @@ ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) { static char filename[1024]; snprintf(filename, sizeof(filename), "%s/ncclUniqueId.sync", fs_path); + if (result->process_rank != 0) { + // Wait for the server to write the file + // This is a naive way to synchronize the processes + sleep(2); + } + if (result->process_rank == 0) { ncclCheck(ncclGetUniqueId(&nccl_id)); idFile = fopen(filename, "wb"); @@ -207,7 +213,7 @@ ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) { } else { // Other ranks wait until the file is available and read the unique ID do { - usleep(1000000); // 1 second + sleep(1); // 1 second idFile = fopen(filename, "rb"); if (idFile != NULL) break; } while (idFile == NULL); diff --git a/scripts/multi_node/run_gpt2_124M_fs.sbatch b/scripts/multi_node/run_gpt2_124M_fs.sbatch new file mode 100755 index 0000000..2498155 --- /dev/null +++ b/scripts/multi_node/run_gpt2_124M_fs.sbatch @@ -0,0 +1,84 @@ +#!/bin/bash +#SBATCH --job-name=llmc-multinode # job name +#SBATCH --output=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.log # output file +#SBATCH --error=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.err # error file +#SBATCH --partition=llmc # Specify the GPU partition +#SBATCH --ntasks=16 # total number of processes to launch on all nodes +#SBATCH --nodes=2 # total number of nodes +#SBATCH --ntasks-per-node=8 # assuming each node has 8 gpus +#SBATCH --gres=gpu:8 # request 8 gpus from each node + +# NOTE: change the above slurm arguments to match your system! +# Run with `sbatch ` + +make train_gpt2cu USE_CUDNN=1 USE_MPI=0 + +# NOTE: change the following to match your system +binary_path="/home/ubuntu/llm.c/train_gpt2cu" +out_dir="/ephemeral/data/fineweb/log_gpt2_124M_multi" +train_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_train_*.bin' +val_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_val_*.bin' +sync_fs_path=$out_dir # needs to be a shared filesystem path that all nodes can access + +# In case the file system is shared this is a no-op. +# Otherwise, we need to copy the binary to all nodes. +current_user=$USER +hosts=$(scontrol show hostnames $SLURM_JOB_NODELIST) # get the hostnames of the allocated nodes +current_host=$(hostname) +for host in $hosts; do + if [ $host == $current_host ]; then + continue + fi + echo "copying $binary_path to $current_user@$host" + scp -r $binary_path $current_user@$host:$binary_path +done + +# Use this for NCCL debugging if you run into issues +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=ALL +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +# Optimization flags +export NCCL_NET_GDR_LEVEL=2 # use GPUDirect RDMA - allows for direct memory access between GPUs across different nodes by bypassing the CPU +export NCCL_IB_DISABLE=0 # use InfiniBand if available + +# NOTE: change the following environment variables to match your system - or comment them out if you don't need them +export NCCL_SOCKET_IFNAME=ens17 +export OMPI_MCA_btl_tcp_if_include=ens17 +export NCCL_P2P_LEVEL=PXB + +if [ -z "$SLURM_JOB_ID" ]; then + echo "Make sure you're running in a SLURM environment. Did you forget to run with sbatch? Aborting." + exit 1 +else + DATESTRING=`date "+%Y-%m-%dT%H:%M:%S"` + echo "Running in a SLURM environment (job ID: $SLURM_JOB_ID, user: $current_user)" + echo "Running on hosts: $(echo $(scontrol show hostname))" + echo "$DATESTRING" +fi + +srun -l -u bash -c " + $binary_path \ + -i '$train_data_path' \ + -j '$val_data_path' \ + -o $out_dir \ + -v 250 -s 20000 -g 144 \ + -h 1 \ + -b 64 -t 1024 \ + -d 2097152 \ + -r 0 \ + -z 1 \ + -c 0.1 \ + -l 0.0006 \ + -q 0.0 \ + -u 700 \ + -n 5000 \ + -y 1 \ + -e d12 \ + -pn \$SLURM_NTASKS \ + -pr \$SLURM_PROCID \ + -pf $sync_fs_path \ + -pi "fs" \ +" + +echo "$DATESTRING" \ No newline at end of file diff --git a/scripts/multi_node/run_gpt2_124M_mpi.sh b/scripts/multi_node/run_gpt2_124M_mpi.sh new file mode 100755 index 0000000..366f8ff --- /dev/null +++ b/scripts/multi_node/run_gpt2_124M_mpi.sh @@ -0,0 +1,49 @@ + +make train_gpt2cu USE_CUDNN=1 USE_MPI=1 + +# NOTE: change the following to match your system +binary_path="/home/ubuntu/llm.c/train_gpt2cu" +out_dir="/ephemeral/data/fineweb/log_gpt2_124M_multi" +train_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_train_*.bin' +val_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_val_*.bin' +# You can find these names either in `/etc/hosts`` file or in the terminal (user@host:~$). +host1="h100-node-1-0" # master and worker node +host2="h100-node-1-1" # worker node + +# In case the file system is shared this is a no-op. +# Otherwise, we need to copy the binary to all nodes. +scp -r $binary_path $USER@$host2:$binary_path + +# Use this for NCCL debugging if you run into issues +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=ALL +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +# Optimization flags +export NCCL_NET_GDR_LEVEL=2 # use GPUDirect RDMA - allows for direct memory access between GPUs across different nodes by bypassing the CPU +export NCCL_IB_DISABLE=0 # use InfiniBand if available + +# NOTE: change the following environment variables to match your system - or comment them out if you don't need them +export NCCL_SOCKET_IFNAME=ens17 +export OMPI_MCA_btl_tcp_if_include=ens17 +export NCCL_P2P_LEVEL=PXB + +mpirun -np 16 --host $host1:8,$host2:8 \ + $binary_path \ + -i "$train_data_path" \ + -j "$val_data_path" \ + -o $out_dir \ + -v 250 -s 20000 -g 144 \ + -h 1 \ + -b 64 -t 1024 \ + -d 2097152 \ + -r 0 \ + -z 1 \ + -c 0.1 \ + -l 0.0006 \ + -q 0.1 \ + -u 700 \ + -n 1000 \ + -y 0 \ + -e d12 \ + -pi "mpi" \ diff --git a/scripts/multi_node/run_gpt2_124M_tcp.sbatch b/scripts/multi_node/run_gpt2_124M_tcp.sbatch new file mode 100755 index 0000000..fbf9c34 --- /dev/null +++ b/scripts/multi_node/run_gpt2_124M_tcp.sbatch @@ -0,0 +1,85 @@ +#!/bin/bash +#SBATCH --job-name=llmc-multinode # job name +#SBATCH --output=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.log # output file +#SBATCH --error=/home/ubuntu/llm.c/scripts/multi_node/%x_%j_%t.err # error file +#SBATCH --partition=llmc # Specify the GPU partition +#SBATCH --ntasks=16 # total number of processes to launch on all nodes +#SBATCH --nodes=2 # total number of nodes +#SBATCH --ntasks-per-node=8 # assuming each node has 8 gpus +#SBATCH --gres=gpu:8 # request 8 gpus from each node + +# NOTE: change the above slurm arguments to match your system! +# Run with `sbatch ` + +make train_gpt2cu USE_CUDNN=1 USE_MPI=0 + +# NOTE: change the following to match your system +binary_path="/home/ubuntu/llm.c/train_gpt2cu" +out_dir="/ephemeral/data/fineweb/log_gpt2_124M_multi" +train_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_train_*.bin' +val_data_path='/ephemeral/data/fineweb/bin_10B/fineweb_val_*.bin' +# NOTE: change the server_ip to the IP address of the machine that is running process zero +server_ip="10.0.1.220" + +# In case the file system is shared this is a no-op. +# Otherwise, we need to copy the binary to all nodes. +current_user=$USER +hosts=$(scontrol show hostnames $SLURM_JOB_NODELIST) # get the hostnames of the allocated nodes +current_host=$(hostname) +for host in $hosts; do + if [ $host == $current_host ]; then + continue + fi + echo "copying $binary_path to $current_user@$host" + scp -r $binary_path $current_user@$host:$binary_path +done + +# Use this for NCCL debugging if you run into issues +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=ALL +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +# Optimization flags +export NCCL_NET_GDR_LEVEL=2 # use GPUDirect RDMA - allows for direct memory access between GPUs across different nodes by bypassing the CPU +export NCCL_IB_DISABLE=0 # use InfiniBand if available + +# NOTE: change the following environment variables to match your system - or comment them out if you don't need them +export NCCL_SOCKET_IFNAME=ens17 +export OMPI_MCA_btl_tcp_if_include=ens17 +export NCCL_P2P_LEVEL=PXB + +if [ -z "$SLURM_JOB_ID" ]; then + echo "Make sure you're running in a SLURM environment. Did you forget to run with sbatch? Aborting." + exit 1 +else + DATESTRING=`date "+%Y-%m-%dT%H:%M:%S"` + echo "Running in a SLURM environment (job ID: $SLURM_JOB_ID, user: $current_user)" + echo "Running on hosts: $(echo $(scontrol show hostname))" + echo "$DATESTRING" +fi + +srun -l -u bash -c " + $binary_path \ + -i '$train_data_path' \ + -j '$val_data_path' \ + -o $out_dir \ + -v 250 -s 20000 -g 144 \ + -h 1 \ + -b 64 -t 1024 \ + -d 2097152 \ + -r 0 \ + -z 1 \ + -c 0.1 \ + -l 0.0006 \ + -q 0.0 \ + -u 700 \ + -n 5000 \ + -y 1 \ + -e d12 \ + -pn \$SLURM_NTASKS \ + -pr \$SLURM_PROCID \ + -ps $server_ip \ + -pi "tcp" \ +" + +echo "$DATESTRING" From 75a2c8b10f705d1ec02fff8fb6c64cb24022be9e Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 11:37:18 +0000 Subject: [PATCH 37/47] Add usage for new args --- train_gpt2.cu | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 215caca..c618046 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1318,6 +1318,12 @@ void error_usage() { // memory management fprintf(stderr, " -z zero_stage, Zero Optimization Stage, 0,1,2,3 (default = 0)\n"); fprintf(stderr, " -r recompute: less memory but less speed. (default = 1), 0|1|2 = none,gelu,gelu+ln\n"); + // multi-node settings + fprintf(stderr, " -pn num_processes (default = 1)\n"); + fprintf(stderr, " -pr process_rank (default = 0)\n"); + fprintf(stderr, " -pm nccl_init_method: tcp,fs,mpi (default = mpi)\n"); + fprintf(stderr, " -ps server_ip - used only when nccl_init_method is tcp (default = -1)\n"); + fprintf(stderr, " -pp fs_path - used only when nccl_init_method is fs (default = /tmp)\n"); exit(EXIT_FAILURE); } @@ -1353,7 +1359,7 @@ int main(int argc, char *argv[]) { // multi-node settings int num_processes = 1; // this should be set by the slurm environment int process_rank = 0; // this should be set by the slurm environment - char nccl_init_method[256] = "tcp"; // "tcp" or "fs" or "mpi" + char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi" char server_ip[256] = "-1"; // used if init_method set to "tcp" -> set to your server ip address char fs_path[256] = "/tmp"; // used if init_method set to "fs" -> set to a shared filesystem path for (int i = 1; i < argc; i+=2) { From 3abdf9a3e4baf05a13b5e86e3aafc85b91e1e9a5 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 12:07:35 +0000 Subject: [PATCH 38/47] Add gpus per node arg --- llmc/zero.cuh | 64 +++++++++++---------- scripts/multi_node/run_gpt2_124M_fs.sbatch | 1 + scripts/multi_node/run_gpt2_124M_tcp.sbatch | 1 + train_gpt2.cu | 5 +- 4 files changed, 40 insertions(+), 31 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index e85e6b8..510b144 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -90,8 +90,8 @@ void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int nu ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) { ncclUniqueId nccl_id; - // hardcode an arbitrary port number between 1024 and 49151 (registered ports) - int SERVER_PORT = 12345; + + int SERVER_PORT = 12345; // hardcoded an arbitrary port number between 1024 and 49151 (registered ports) if (result->process_rank == 0) { ncclCheck(ncclGetUniqueId(&nccl_id)); @@ -103,13 +103,13 @@ ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) int addrlen = sizeof(address); int opt = 1; - // Create a TCP socket + // Step 1) create a server TCP socket if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { printf("Socket failed"); exit(EXIT_FAILURE); } - // set socket options + // Step 2) set socket options // SOL_SOCKET - means that option is configured at socket level // SO_REUSEADDR - allows to bind to an address which is in a TIME_WAIT state (already used by another socket) - useful when restarting the server // SO_REUSEPORT - allows to bind to the same port multiple times @@ -118,22 +118,24 @@ ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) exit(EXIT_FAILURE); } + // Step 3) set the server address and port address.sin_family = AF_INET; // IPv4 address.sin_addr.s_addr = inet_addr(server_ip); // alternatively use INADDR_ANY to bind to all interfaces, currently we only allow ethernet address.sin_port = htons(SERVER_PORT); - // Bind the socket to the address and port + // Step 4) bind the socket to the address and port if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) < 0) { printf("Bind failed"); exit(EXIT_FAILURE); } - // MAX_CLIENTS specifies the maximum number of clients that can be queued for this server + // Step 5) MAX_CLIENTS specifies the maximum number of clients that can be queued for this server if (listen(server_socket, MAX_CLIENTS) < 0) { printf("Listen failed"); exit(EXIT_FAILURE); } + // Step 6) accept connections from clients printf("Waiting for clients to connect...\n"); while (num_clients < MAX_CLIENTS) { if ((new_socket = accept(server_socket, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) { @@ -144,24 +146,24 @@ ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) printf("Client %d connected\n", num_clients); } + // Step 7) send the NCCL ID to all clients send_nccl_id_to_clients(&nccl_id, client_sockets, num_clients); printf("NCCL ID sent to all clients\n"); close(server_socket); } else { - int num_attempts = 5; + int num_connection_attempts = 5; int time_to_sleep = 2; - int client_socket; struct sockaddr_in serv_addr; - // Create a TCP socket + // Step 1) create a client TCP socket if ((client_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { printf("Socket creation error"); exit(EXIT_FAILURE); } - // Set the server address and port + // Step 2) set the server address and port serv_addr.sin_family = AF_INET; serv_addr.sin_port = htons(SERVER_PORT); if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) { @@ -169,16 +171,17 @@ ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) exit(EXIT_FAILURE); } - // Try to connect to the server - retry if connection fails + // Step 3) Try to connect to the server - retry up to `num_connection_attempts` times if the connection fails while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) { printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep); - if (--num_attempts == 0) { + if (--num_connection_attempts == 0) { printf("Failed to connect to the server\n"); exit(EXIT_FAILURE); } sleep(time_to_sleep); } + // Step 4) receive the NCCL ID from the server if (recv(client_socket, &nccl_id, sizeof(nccl_id), 0) <= 0) { printf("Failed to receive nccl_id"); exit(EXIT_FAILURE); @@ -198,9 +201,8 @@ ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) { static char filename[1024]; snprintf(filename, sizeof(filename), "%s/ncclUniqueId.sync", fs_path); - if (result->process_rank != 0) { - // Wait for the server to write the file - // This is a naive way to synchronize the processes + if (result->process_rank != 0) { // client processse should wait for the server to write to the file + // This is a naive and not 100% robust way to synchronize the processes but it should work almost always sleep(2); } @@ -267,23 +269,13 @@ int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { #endif -MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank, char* server_ip, char* fs_path, char* init_method) { +MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank, int gpus_per_node, char* server_ip, char* fs_path, char* init_method) { #ifdef MULTI_GPU MultiGpuConfig result; ncclUniqueId nccl_id; - if (strcmp(init_method, "mpi") != 0) { - result.process_rank = process_rank; - result.num_processes = num_processes; - result.local_device_idx = process_rank % 8; - if (strcmp(init_method, "tcp") == 0) { - nccl_id = get_nccl_id_via_tcp(&result, server_ip); - } else if (strcmp(init_method, "fs") == 0) { - nccl_id = get_nccl_id_via_fs(&result, fs_path); - } else { - printf("Invalid init method\n"); - exit(EXIT_FAILURE); - } - } else { + // Get nccl_id using MPI, TCP, or FS (file system synchronization) methods + // On newer slurm versions (slurm-wlm package) PMIx is disabled so we can not use MPI for NCCL init in multi node setup + if (strcmp(init_method, "mpi") == 0) { #ifdef USE_MPI mpiCheck(MPI_Init(NULL, NULL)); mpiCheck(MPI_Comm_rank(MPI_COMM_WORLD, &result.process_rank)); @@ -294,9 +286,21 @@ MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank, char* } mpiCheck(MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD)); #else - printf("MPI support is disabled. Please enable MPI support to use MPI init method.\n"); + printf("MPI support is disabled. Please enable MPI support to use MPI-based NCCL-init method.\n"); exit(EXIT_FAILURE); #endif + } else { + result.process_rank = process_rank; + result.num_processes = num_processes; + result.local_device_idx = process_rank % gpus_per_node; + if (strcmp(init_method, "tcp") == 0) { + nccl_id = get_nccl_id_via_tcp(&result, server_ip); + } else if (strcmp(init_method, "fs") == 0) { + nccl_id = get_nccl_id_via_fs(&result, fs_path); + } else { + printf("Invalid NCCL-init method\n"); + exit(EXIT_FAILURE); + } } cudaCheck(cudaSetDevice(result.local_device_idx)); ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank)); diff --git a/scripts/multi_node/run_gpt2_124M_fs.sbatch b/scripts/multi_node/run_gpt2_124M_fs.sbatch index 2498155..d2e4b5f 100755 --- a/scripts/multi_node/run_gpt2_124M_fs.sbatch +++ b/scripts/multi_node/run_gpt2_124M_fs.sbatch @@ -77,6 +77,7 @@ srun -l -u bash -c " -e d12 \ -pn \$SLURM_NTASKS \ -pr \$SLURM_PROCID \ + -pg \$SLURM_NTASKS_PER_NODE \ -pf $sync_fs_path \ -pi "fs" \ " diff --git a/scripts/multi_node/run_gpt2_124M_tcp.sbatch b/scripts/multi_node/run_gpt2_124M_tcp.sbatch index fbf9c34..416dfb8 100755 --- a/scripts/multi_node/run_gpt2_124M_tcp.sbatch +++ b/scripts/multi_node/run_gpt2_124M_tcp.sbatch @@ -78,6 +78,7 @@ srun -l -u bash -c " -e d12 \ -pn \$SLURM_NTASKS \ -pr \$SLURM_PROCID \ + -pg \$SLURM_NTASKS_PER_NODE \ -ps $server_ip \ -pi "tcp" \ " diff --git a/train_gpt2.cu b/train_gpt2.cu index c618046..97b0b53 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1321,6 +1321,7 @@ void error_usage() { // multi-node settings fprintf(stderr, " -pn num_processes (default = 1)\n"); fprintf(stderr, " -pr process_rank (default = 0)\n"); + fprintf(stderr, " -pg gpus_per_node (default = 8)\n"); fprintf(stderr, " -pm nccl_init_method: tcp,fs,mpi (default = mpi)\n"); fprintf(stderr, " -ps server_ip - used only when nccl_init_method is tcp (default = -1)\n"); fprintf(stderr, " -pp fs_path - used only when nccl_init_method is fs (default = /tmp)\n"); @@ -1359,6 +1360,7 @@ int main(int argc, char *argv[]) { // multi-node settings int num_processes = 1; // this should be set by the slurm environment int process_rank = 0; // this should be set by the slurm environment + int gpus_per_node = 8; // this should be set by the slurm environment char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi" char server_ip[256] = "-1"; // used if init_method set to "tcp" -> set to your server ip address char fs_path[256] = "/tmp"; // used if init_method set to "fs" -> set to a shared filesystem path @@ -1397,9 +1399,10 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 'p' && argv[i][2] == 's') { strcpy(server_ip, argv[i+1]); } else if (argv[i][1] == 'p' && argv[i][2] == 'n') { num_processes = atoi(argv[i+1]); } else if (argv[i][1] == 'p' && argv[i][2] == 'r') { process_rank = atoi(argv[i+1]); } + else if (argv[i][1] == 'p' && argv[i][2] == 'g') { gpus_per_node = atoi(argv[i+1]); } else { error_usage(); } } - multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, server_ip, fs_path, nccl_init_method); + multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method); // should do a bit more error checking here assert(warmup_iterations >= 0); From a259eeabe5da05d3a09c6aab91c5cf94da6e786c Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 14:09:42 +0000 Subject: [PATCH 39/47] Fix CI errors --- Makefile | 6 +++--- profile_gpt2.cu | 8 +++++++- scripts/multi_node/run_gpt2_124M_fs.sbatch | 2 +- scripts/multi_node/run_gpt2_124M_mpi.sh | 2 +- scripts/multi_node/run_gpt2_124M_tcp.sbatch | 2 +- test_gpt2.cu | 8 +++++++- 6 files changed, 20 insertions(+), 8 deletions(-) diff --git a/Makefile b/Makefile index 200b81c..c768140 100644 --- a/Makefile +++ b/Makefile @@ -207,14 +207,14 @@ else endif endif -ifeq ($(USE_MPI), 1) +ifeq ($(NO_USE_MPI), 1) + $(info → MPI is manually disabled) +else $(info → MPI is manually enabled) NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/ NVCC_FLAGS += -DUSE_MPI NVCC_LDLIBS += -lmpi -else - $(info → MPI is manually disabled) endif # Precision settings, default to bf16 but ability to override diff --git a/profile_gpt2.cu b/profile_gpt2.cu index 8a7ddbf..ab1c4e0 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -28,7 +28,13 @@ the profile.ncu-rep from a cloud box to local to pretty view. #include "train_gpt2.cu" int main(int argc, char *argv[]) { - multi_gpu_config = multi_gpu_config_init(&argc, &argv); + char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi" + int num_processes = -1; // doesn't matter when using MPI + int process_rank = -1; // doesn't matter when using MPI + int gpus_per_node = -1; // doesn't matter when using MPI + char server_ip[256] = ""; // doesn't matter when using MPI + char fs_path[256] = ""; // doesn't matter when using MPI + multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, nccl_init_method, server_ip, fs_path); common_start(true, true); // build the GPT-2 model from a checkpoint diff --git a/scripts/multi_node/run_gpt2_124M_fs.sbatch b/scripts/multi_node/run_gpt2_124M_fs.sbatch index d2e4b5f..9bef9aa 100755 --- a/scripts/multi_node/run_gpt2_124M_fs.sbatch +++ b/scripts/multi_node/run_gpt2_124M_fs.sbatch @@ -11,7 +11,7 @@ # NOTE: change the above slurm arguments to match your system! # Run with `sbatch ` -make train_gpt2cu USE_CUDNN=1 USE_MPI=0 +make train_gpt2cu USE_CUDNN=1 NO_USE_MPI=1 # NOTE: change the following to match your system binary_path="/home/ubuntu/llm.c/train_gpt2cu" diff --git a/scripts/multi_node/run_gpt2_124M_mpi.sh b/scripts/multi_node/run_gpt2_124M_mpi.sh index 366f8ff..e09b027 100755 --- a/scripts/multi_node/run_gpt2_124M_mpi.sh +++ b/scripts/multi_node/run_gpt2_124M_mpi.sh @@ -1,5 +1,5 @@ -make train_gpt2cu USE_CUDNN=1 USE_MPI=1 +make train_gpt2cu USE_CUDNN=1 # NOTE: change the following to match your system binary_path="/home/ubuntu/llm.c/train_gpt2cu" diff --git a/scripts/multi_node/run_gpt2_124M_tcp.sbatch b/scripts/multi_node/run_gpt2_124M_tcp.sbatch index 416dfb8..f6cd3a7 100755 --- a/scripts/multi_node/run_gpt2_124M_tcp.sbatch +++ b/scripts/multi_node/run_gpt2_124M_tcp.sbatch @@ -11,7 +11,7 @@ # NOTE: change the above slurm arguments to match your system! # Run with `sbatch ` -make train_gpt2cu USE_CUDNN=1 USE_MPI=0 +make train_gpt2cu USE_CUDNN=1 NO_USE_MPI=1 # NOTE: change the following to match your system binary_path="/home/ubuntu/llm.c/train_gpt2cu" diff --git a/test_gpt2.cu b/test_gpt2.cu index eee87dd..854c013 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -89,7 +89,13 @@ float* float_cpu_malloc_and_point_parameters(FloatParameterTensors* params, size } int main(int argc, char *argv[]) { - multi_gpu_config = multi_gpu_config_init(&argc, &argv); + char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi" + int num_processes = -1; // doesn't matter when using MPI + int process_rank = -1; // doesn't matter when using MPI + int gpus_per_node = -1; // doesn't matter when using MPI + char server_ip[256] = ""; // doesn't matter when using MPI + char fs_path[256] = ""; // doesn't matter when using MPI + multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, nccl_init_method, server_ip, fs_path); common_start(false, true); // set the right paths From fbf4f59c8fca5d71b5273ad851cb5f41c510fb87 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 14:25:01 +0000 Subject: [PATCH 40/47] Add MPI under multi gpu guard --- llmc/zero.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index 510b144..a764143 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -14,11 +14,10 @@ Utilities for ZeRO sharding #ifdef MULTI_GPU #include -#endif - #ifdef USE_MPI #include #endif +#endif // ---------------------------------------------------------------------------- // Multi-GPU related From c420a522b25e5f21fac3e953b31719b31e9661a6 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 14:38:49 +0000 Subject: [PATCH 41/47] Fix arg permutation --- Makefile | 2 +- profile_gpt2.cu | 2 +- test_gpt2.cu | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index c768140..6899ed7 100644 --- a/Makefile +++ b/Makefile @@ -210,7 +210,7 @@ endif ifeq ($(NO_USE_MPI), 1) $(info → MPI is manually disabled) else - $(info → MPI is manually enabled) + $(info ✓ MPI enabled) NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/ NVCC_FLAGS += -DUSE_MPI diff --git a/profile_gpt2.cu b/profile_gpt2.cu index ab1c4e0..62f9354 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -34,7 +34,7 @@ int main(int argc, char *argv[]) { int gpus_per_node = -1; // doesn't matter when using MPI char server_ip[256] = ""; // doesn't matter when using MPI char fs_path[256] = ""; // doesn't matter when using MPI - multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, nccl_init_method, server_ip, fs_path); + multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method); common_start(true, true); // build the GPT-2 model from a checkpoint diff --git a/test_gpt2.cu b/test_gpt2.cu index 854c013..45964d5 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -95,7 +95,7 @@ int main(int argc, char *argv[]) { int gpus_per_node = -1; // doesn't matter when using MPI char server_ip[256] = ""; // doesn't matter when using MPI char fs_path[256] = ""; // doesn't matter when using MPI - multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, nccl_init_method, server_ip, fs_path); + multi_gpu_config = multi_gpu_config_init(num_processes, process_rank, gpus_per_node, server_ip, fs_path, nccl_init_method); common_start(false, true); // set the right paths From 05ef1e7dc0c6275c44a88cb7669d6efc22580efb Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 14:53:58 +0000 Subject: [PATCH 42/47] Fix makefile --- Makefile | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 6899ed7..1a2fa02 100644 --- a/Makefile +++ b/Makefile @@ -199,7 +199,6 @@ else else ifeq ($(shell dpkg -l | grep -q nccl && echo "exists"), exists) $(info ✓ NCCL found, OK to train with multiple GPUs) NVCC_LDLIBS += -lnccl - NVCC_FLAGS += -DMULTI_GPU else $(info ✗ NCCL is not found, disabling multi-GPU support) $(info ---> On Linux you can try install NCCL with `sudo apt install libnccl2 libnccl-dev`) @@ -209,12 +208,15 @@ endif ifeq ($(NO_USE_MPI), 1) $(info → MPI is manually disabled) -else +else ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists) $(info ✓ MPI enabled) NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/ - NVCC_FLAGS += -DUSE_MPI NVCC_LDLIBS += -lmpi + NVCC_FLAGS += -DUSE_MPI + NVCC_FLAGS += -DMULTI_GPU +else + $(info ✗ MPI not found) endif # Precision settings, default to bf16 but ability to override From f9d0e8b67ada1ad93437389354b0f983c084f697 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 18:26:53 +0000 Subject: [PATCH 43/47] Add socket logic for Windows --- Makefile | 2 +- llmc/zero.cuh | 143 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 1a2fa02..b8326e3 100644 --- a/Makefile +++ b/Makefile @@ -88,7 +88,7 @@ else LDFLAGS := LDLIBS := INCLUDES := - NVCC_FLAGS += -I"dev" + NVCC_FLAGS += -I"dev" -lWs2_32 ifeq ($(WIN_CI_BUILD),1) $(info Windows CI build) OUTPUT_FILE = /link /OUT:$@ diff --git a/llmc/zero.cuh b/llmc/zero.cuh index a764143..df2df6a 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -5,7 +5,13 @@ Utilities for ZeRO sharding #ifndef LLMC_ZERO_CUH #define LLMC_ZERO_CUH +#ifdef _WIN32 +#include +#include +#else #include +#endif + #include #include #include @@ -87,6 +93,19 @@ void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int nu } } +#ifdef _WIN32 +void send_nccl_id_to_clients_windows(ncclUniqueId *nccl_id, SOCKET client_sockets[], int num_clients) { + for (int i = 0; i < num_clients; ++i) { + if (send(client_sockets[i], (const char *)nccl_id, sizeof(*nccl_id), 0) == SOCKET_ERROR) { + printf("Failed to send nccl_id"); + WSACleanup(); + exit(EXIT_FAILURE); + } + closesocket(client_sockets[i]); + } +} +#endif + ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) { ncclUniqueId nccl_id; @@ -193,6 +212,126 @@ ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) return nccl_id; } +#ifdef _WIN32 +// Same as get_nccl_id_via_tcp but for Windows +ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* server_ip) { + ncclUniqueId nccl_id; + + int SERVER_PORT = 12345; // hardcoded an arbitrary port number between 1024 and 49151 (registered ports) + WSADATA wsaData; + if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) { + printf("WSAStartup failed"); + exit(EXIT_FAILURE); + } + + if (result->process_rank == 0) { + ncclCheck(ncclGetUniqueId(&nccl_id)); + + int MAX_CLIENTS = result->num_processes - 1; + SOCKET client_sockets[MAX_CLIENTS]; + int num_clients = 0; + SOCKET server_socket, new_socket; + struct sockaddr_in address; + int addrlen = sizeof(address); + + // Step 1) create a server TCP socket + if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) { + printf("Socket failed"); + WSACleanup(); + exit(EXIT_FAILURE); + } + + // Step 2) set the server address and port + address.sin_family = AF_INET; // IPv4 + address.sin_addr.s_addr = inet_addr(server_ip); + address.sin_port = htons(SERVER_PORT); + + // Step 3) bind the socket to the address and port + if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) == SOCKET_ERROR) { + printf("Bind failed"); + closesocket(server_socket); + WSACleanup(); + exit(EXIT_FAILURE); + } + + // Step 4) MAX_CLIENTS specifies the maximum number of clients that can be queued for this server + if (listen(server_socket, MAX_CLIENTS) == SOCKET_ERROR) { + printf("Listen failed"); + closesocket(server_socket); + WSACleanup(); + exit(EXIT_FAILURE); + } + + // Step 5) accept connections from clients + printf("Waiting for clients to connect...\n"); + while (num_clients < MAX_CLIENTS) { + if ((new_socket = accept(server_socket, (struct sockaddr *)&address, &addrlen)) == INVALID_SOCKET) { + printf("Accept failed"); + closesocket(server_socket); + WSACleanup(); + exit(EXIT_FAILURE); + } + client_sockets[num_clients++] = new_socket; + printf("Client %d connected\n", num_clients); + } + + // Step 6) send the NCCL ID to all clients + send_nccl_id_to_clients_windows(&nccl_id, client_sockets, num_clients); + printf("NCCL ID sent to all clients\n"); + + closesocket(server_socket); + } else { + int num_connection_attempts = 5; + int time_to_sleep = 2; + SOCKET client_socket; + struct sockaddr_in serv_addr; + + // Step 1) create a client TCP socket + if ((client_socket = socket(AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET) { + printf("Socket creation error"); + WSACleanup(); + exit(EXIT_FAILURE); + } + + // Step 2) set the server address and port + serv_addr.sin_family = AF_INET; + serv_addr.sin_port = htons(SERVER_PORT); + if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) { + printf("Invalid address or address not supported"); + closesocket(client_socket); + WSACleanup(); + exit(EXIT_FAILURE); + } + + // Step 3) Try to connect to the server - retry up to `num_connection_attempts` times if the connection fails + while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) == SOCKET_ERROR) { + printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep); + if (--num_connection_attempts == 0) { + printf("Failed to connect to the server\n"); + closesocket(client_socket); + WSACleanup(); + exit(EXIT_FAILURE); + } + Sleep(time_to_sleep * 1000); + } + + // Step 4) receive the NCCL ID from the server + if (recv(client_socket, (char *)&nccl_id, sizeof(nccl_id), 0) <= 0) { + printf("Failed to receive nccl_id"); + closesocket(client_socket); + WSACleanup(); + exit(EXIT_FAILURE); + } + + printf("Received NCCL ID\n"); + closesocket(client_socket); + } + + WSACleanup(); + return nccl_id; +} +#endif + ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) { // Works assuming that the filesystem is shared among all processes ncclUniqueId nccl_id; @@ -293,7 +432,11 @@ MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank, int gp result.num_processes = num_processes; result.local_device_idx = process_rank % gpus_per_node; if (strcmp(init_method, "tcp") == 0) { + #ifdef _WIN32 + nccl_id = get_nccl_id_via_tcp_windows(&result, server_ip); + #else nccl_id = get_nccl_id_via_tcp(&result, server_ip); + #endif } else if (strcmp(init_method, "fs") == 0) { nccl_id = get_nccl_id_via_fs(&result, fs_path); } else { From 60c95c82aae1720aa65926ee51f40bf19d5ba02b Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 18:50:56 +0000 Subject: [PATCH 44/47] Add a multi node readme section --- README.md | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index fe24a72..82bd2fe 100644 --- a/README.md +++ b/README.md @@ -134,7 +134,9 @@ sudo apt-get -y install libcudnn9-dev-cuda-12 On top of this you need the [cuDNN frontend](https://github.com/NVIDIA/cudnn-frontend/tree/main), but this is just header files. Simply clone the repo to your disk. The Makefile currently looks for it in either your home directory or the current directory. If you have put it elsewhere, add `CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include` to the `make` command-line. -**multi-GPU training using MPI and NCCL**. Make sure you install MPI and NCCL, e.g. on Linux: +## multi-GPU training + +Make sure you install MPI and NCCL, e.g. on Linux: ```bash sudo apt install openmpi-bin openmpi-doc libopenmpi-dev @@ -149,6 +151,23 @@ make train_gpt2cu mpirun -np ./train_gpt2cu ``` +or simply run one of our scripts under `./scripts/`. + +## multi-node training + +Make sure you've installed `NCCL` following instructions from [multi-GPU](#multi-gpu-training) section. + +There are 3 ways we currently support that allow you to run multi-node training: +1) Use OpenMPI to exchange nccl id and initialize NCCL. See e.g. `./scripts/multi_node/run_gpt2_124M_mpi.sh` script for details. +2) Use shared file system to init NCCL. See `./scripts/multi_node/run_gpt2_124M_fs.sbatch` script for details. +3) Use TCP sockets to init NCCL. See `./scripts/multi_node/run_gpt2_124M_tcp.sbatch` script for details. + +Note: +* If you're running in a slurm environment and your slurm doesn't support PMIx (which we assume will be a common situation given that `slurm-wlm` dropped PMIx support) you will have to use FS (2) or TCP (3) approach. To test whether your slurm supports PMIx run: `srun --mpi=list` and see whether you get `pmix` in the output. +* If you don't have slurm set up, you can kick off a multi-node run using `mpirun` - MPI (1). + +None of these 3 methods is superior, we just offer you options so that you can run in your specific environment. + ## experiments / sweeps Just as an example process to sweep learning rates on a machine with 4 GPUs on TinyStories. Run a shell script `sweep.sh` (after you of course `chmod u+x sweep.sh`): From ac1e37b6593dd8f894d519c8dca956fce871e2b0 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 18:56:48 +0000 Subject: [PATCH 45/47] Remove linux socket funcs when windows --- llmc/zero.cuh | 225 +++++++++++++++++++++++++------------------------- 1 file changed, 113 insertions(+), 112 deletions(-) diff --git a/llmc/zero.cuh b/llmc/zero.cuh index df2df6a..a46ed75 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -83,15 +83,6 @@ typedef struct { } MultiGpuConfig; #ifdef MULTI_GPU -void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int num_clients) { - for (int i = 0; i < num_clients; ++i) { - if (send(client_sockets[i], nccl_id, sizeof(*nccl_id), 0) == -1) { - printf("Failed to send nccl_id"); - exit(EXIT_FAILURE); - } - close(client_sockets[i]); - } -} #ifdef _WIN32 void send_nccl_id_to_clients_windows(ncclUniqueId *nccl_id, SOCKET client_sockets[], int num_clients) { @@ -104,113 +95,17 @@ void send_nccl_id_to_clients_windows(ncclUniqueId *nccl_id, SOCKET client_socket closesocket(client_sockets[i]); } } -#endif - -ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) { - ncclUniqueId nccl_id; - - int SERVER_PORT = 12345; // hardcoded an arbitrary port number between 1024 and 49151 (registered ports) - if (result->process_rank == 0) { - ncclCheck(ncclGetUniqueId(&nccl_id)); - - int MAX_CLIENTS = result->num_processes - 1; - int client_sockets[MAX_CLIENTS]; - int num_clients = 0; - int server_socket, new_socket; - struct sockaddr_in address; - int addrlen = sizeof(address); - int opt = 1; - - // Step 1) create a server TCP socket - if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { - printf("Socket failed"); +#else +void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int num_clients) { + for (int i = 0; i < num_clients; ++i) { + if (send(client_sockets[i], nccl_id, sizeof(*nccl_id), 0) == -1) { + printf("Failed to send nccl_id"); exit(EXIT_FAILURE); } - - // Step 2) set socket options - // SOL_SOCKET - means that option is configured at socket level - // SO_REUSEADDR - allows to bind to an address which is in a TIME_WAIT state (already used by another socket) - useful when restarting the server - // SO_REUSEPORT - allows to bind to the same port multiple times - if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)) < 0) { - printf("Setsockopt failed"); - exit(EXIT_FAILURE); - } - - // Step 3) set the server address and port - address.sin_family = AF_INET; // IPv4 - address.sin_addr.s_addr = inet_addr(server_ip); // alternatively use INADDR_ANY to bind to all interfaces, currently we only allow ethernet - address.sin_port = htons(SERVER_PORT); - - // Step 4) bind the socket to the address and port - if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) < 0) { - printf("Bind failed"); - exit(EXIT_FAILURE); - } - - // Step 5) MAX_CLIENTS specifies the maximum number of clients that can be queued for this server - if (listen(server_socket, MAX_CLIENTS) < 0) { - printf("Listen failed"); - exit(EXIT_FAILURE); - } - - // Step 6) accept connections from clients - printf("Waiting for clients to connect...\n"); - while (num_clients < MAX_CLIENTS) { - if ((new_socket = accept(server_socket, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) { - printf("Accept failed"); - exit(EXIT_FAILURE); - } - client_sockets[num_clients++] = new_socket; - printf("Client %d connected\n", num_clients); - } - - // Step 7) send the NCCL ID to all clients - send_nccl_id_to_clients(&nccl_id, client_sockets, num_clients); - printf("NCCL ID sent to all clients\n"); - - close(server_socket); - } else { - int num_connection_attempts = 5; - int time_to_sleep = 2; - int client_socket; - struct sockaddr_in serv_addr; - - // Step 1) create a client TCP socket - if ((client_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { - printf("Socket creation error"); - exit(EXIT_FAILURE); - } - - // Step 2) set the server address and port - serv_addr.sin_family = AF_INET; - serv_addr.sin_port = htons(SERVER_PORT); - if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) { - printf("Invalid address or address not supported"); - exit(EXIT_FAILURE); - } - - // Step 3) Try to connect to the server - retry up to `num_connection_attempts` times if the connection fails - while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) { - printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep); - if (--num_connection_attempts == 0) { - printf("Failed to connect to the server\n"); - exit(EXIT_FAILURE); - } - sleep(time_to_sleep); - } - - // Step 4) receive the NCCL ID from the server - if (recv(client_socket, &nccl_id, sizeof(nccl_id), 0) <= 0) { - printf("Failed to receive nccl_id"); - exit(EXIT_FAILURE); - } - - printf("Received NCCL ID\n"); - close(client_socket); + close(client_sockets[i]); } - - return nccl_id; } +#endif #ifdef _WIN32 // Same as get_nccl_id_via_tcp but for Windows @@ -330,6 +225,112 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser WSACleanup(); return nccl_id; } +#else +ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip) { + ncclUniqueId nccl_id; + + int SERVER_PORT = 12345; // hardcoded an arbitrary port number between 1024 and 49151 (registered ports) + if (result->process_rank == 0) { + ncclCheck(ncclGetUniqueId(&nccl_id)); + + int MAX_CLIENTS = result->num_processes - 1; + int client_sockets[MAX_CLIENTS]; + int num_clients = 0; + int server_socket, new_socket; + struct sockaddr_in address; + int addrlen = sizeof(address); + int opt = 1; + + // Step 1) create a server TCP socket + if ((server_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { + printf("Socket failed"); + exit(EXIT_FAILURE); + } + + // Step 2) set socket options + // SOL_SOCKET - means that option is configured at socket level + // SO_REUSEADDR - allows to bind to an address which is in a TIME_WAIT state (already used by another socket) - useful when restarting the server + // SO_REUSEPORT - allows to bind to the same port multiple times + if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)) < 0) { + printf("Setsockopt failed"); + exit(EXIT_FAILURE); + } + + // Step 3) set the server address and port + address.sin_family = AF_INET; // IPv4 + address.sin_addr.s_addr = inet_addr(server_ip); // alternatively use INADDR_ANY to bind to all interfaces, currently we only allow ethernet + address.sin_port = htons(SERVER_PORT); + + // Step 4) bind the socket to the address and port + if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) < 0) { + printf("Bind failed"); + exit(EXIT_FAILURE); + } + + // Step 5) MAX_CLIENTS specifies the maximum number of clients that can be queued for this server + if (listen(server_socket, MAX_CLIENTS) < 0) { + printf("Listen failed"); + exit(EXIT_FAILURE); + } + + // Step 6) accept connections from clients + printf("Waiting for clients to connect...\n"); + while (num_clients < MAX_CLIENTS) { + if ((new_socket = accept(server_socket, (struct sockaddr *)&address, (socklen_t*)&addrlen)) < 0) { + printf("Accept failed"); + exit(EXIT_FAILURE); + } + client_sockets[num_clients++] = new_socket; + printf("Client %d connected\n", num_clients); + } + + // Step 7) send the NCCL ID to all clients + send_nccl_id_to_clients(&nccl_id, client_sockets, num_clients); + printf("NCCL ID sent to all clients\n"); + + close(server_socket); + } else { + int num_connection_attempts = 5; + int time_to_sleep = 2; + int client_socket; + struct sockaddr_in serv_addr; + + // Step 1) create a client TCP socket + if ((client_socket = socket(AF_INET, SOCK_STREAM, 0)) < 0) { + printf("Socket creation error"); + exit(EXIT_FAILURE); + } + + // Step 2) set the server address and port + serv_addr.sin_family = AF_INET; + serv_addr.sin_port = htons(SERVER_PORT); + if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) { + printf("Invalid address or address not supported"); + exit(EXIT_FAILURE); + } + + // Step 3) Try to connect to the server - retry up to `num_connection_attempts` times if the connection fails + while (connect(client_socket, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) { + printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep); + if (--num_connection_attempts == 0) { + printf("Failed to connect to the server\n"); + exit(EXIT_FAILURE); + } + sleep(time_to_sleep); + } + + // Step 4) receive the NCCL ID from the server + if (recv(client_socket, &nccl_id, sizeof(nccl_id), 0) <= 0) { + printf("Failed to receive nccl_id"); + exit(EXIT_FAILURE); + } + + printf("Received NCCL ID\n"); + close(client_socket); + } + + return nccl_id; +} #endif ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) { From 2667aa1e3eea05bfd719221bbaed89d4a7120eb6 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 19:12:50 +0000 Subject: [PATCH 46/47] Fix windows include issue --- Makefile | 2 +- llmc/zero.cuh | 1 - train_gpt2.cu | 4 ++++ 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index b8326e3..1a2fa02 100644 --- a/Makefile +++ b/Makefile @@ -88,7 +88,7 @@ else LDFLAGS := LDLIBS := INCLUDES := - NVCC_FLAGS += -I"dev" -lWs2_32 + NVCC_FLAGS += -I"dev" ifeq ($(WIN_CI_BUILD),1) $(info Windows CI build) OUTPUT_FILE = /link /OUT:$@ diff --git a/llmc/zero.cuh b/llmc/zero.cuh index a46ed75..cb0faf8 100644 --- a/llmc/zero.cuh +++ b/llmc/zero.cuh @@ -7,7 +7,6 @@ Utilities for ZeRO sharding #ifdef _WIN32 #include -#include #else #include #endif diff --git a/train_gpt2.cu b/train_gpt2.cu index 97b0b53..8e2436d 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2,6 +2,10 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. */ +#ifdef _WIN32 +#define WIN32_LEAN_AND_MEAN +#endif + #include #include #include From 4af6a6ab3d72972bf9b4922874ef8f3a81202e67 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Mon, 24 Jun 2024 21:26:55 +0000 Subject: [PATCH 47/47] Remove /tmp as default init value for fs_path --- train_gpt2.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 8e2436d..43f5617 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1366,8 +1366,8 @@ int main(int argc, char *argv[]) { int process_rank = 0; // this should be set by the slurm environment int gpus_per_node = 8; // this should be set by the slurm environment char nccl_init_method[256] = "mpi"; // "tcp" or "fs" or "mpi" - char server_ip[256] = "-1"; // used if init_method set to "tcp" -> set to your server ip address - char fs_path[256] = "/tmp"; // used if init_method set to "fs" -> set to a shared filesystem path + char server_ip[256] = ""; // used if init_method set to "tcp" -> set to your server ip address + char fs_path[256] = ""; // used if init_method set to "fs" -> set to a shared filesystem path 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