From a00a3a88ef9b6485541fe395427a54c0800cd813 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 26 Apr 2024 18:11:09 +0000 Subject: [PATCH] fix bug with all_hostsname_hashes missing, re-arrange code a bit, make sure we only print on rank0 --- Makefile | 2 +- dev/cuda/nccl_all_reduce.cu | 1 + train_gpt2.cu | 284 +++++++++++++++++++----------------- train_gpt2_fp32.cu | 1 + 4 files changed, 152 insertions(+), 136 deletions(-) diff --git a/Makefile b/Makefile index a2daf69..2908e27 100644 --- a/Makefile +++ b/Makefile @@ -74,7 +74,7 @@ else ifeq ($(shell uname), Darwin) $(warning 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 Adding OpenMPI support) + $(info OpenMPI found, adding support) 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 diff --git a/dev/cuda/nccl_all_reduce.cu b/dev/cuda/nccl_all_reduce.cu index 7559c5b..3bc9564 100644 --- a/dev/cuda/nccl_all_reduce.cu +++ b/dev/cuda/nccl_all_reduce.cu @@ -83,6 +83,7 @@ int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { // 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. diff --git a/train_gpt2.cu b/train_gpt2.cu index e9a7834..83108d3 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -18,6 +18,7 @@ and fp8 (coming soon^TM). #include #include +#include #include #include #include @@ -315,6 +316,106 @@ void *malloc_check(size_t size, const char *file, int line) { #define mallocCheck(size) malloc_check(size, __FILE__, __LINE__) +// ---------------------------------------------------------------------------- +// 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 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. +#ifdef MULTI_GPU + ncclComm_t nccl_comm; // NCCL communication primitive, used for collective mutli-GPU work. +#endif +} MultiGpuConfig; + +// one global variable to hold the multi-GPU configuration for this process +MultiGpuConfig multi_gpu_config; + +#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 = 5381; + for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5) + 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); + cudaCheck(cudaSetDevice(result.local_device_idx)); + ncclUniqueId nccl_id; + if (result.process_rank == 0) { + ncclCheck(ncclGetUniqueId(&nccl_id)); + } + 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)); + return result; +#else + printf("Multi-GPU support is disabled. Using a single GPU."); + return MultiGpuConfig{ + .process_rank = 0, + .num_processes = 1, + .local_device_idx = 0, + }; +#endif +} + +void multi_gpu_config_free(const MultiGpuConfig* multi_gpu_config) { +#ifdef MULTI_GPU + ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm)); + mpiCheck(MPI_Finalize()); +#endif +} + +// convenience function that only prints if the rank of process is zero +void printf0(const char *format, ...) { + if (multi_gpu_config.process_rank == 0) { + va_list args; + va_start(args, format); + vprintf(format, args); + va_end(args); + } +} + // ---------------------------------------------------------------------------- // all the kernels @@ -1564,89 +1665,6 @@ void* malloc_and_point_backward(GradActTensors* acts, const size_t* act_sizes) { return malloc_and_point(ptrs, act_sizes, NUM_BACKWARD_TENSORS); } -// 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 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. -#ifdef MULTI_GPU - ncclComm_t nccl_comm; // NCCL communication primitive, used for collective mutli-GPU work. -#endif -} MultiGpuConfig; - -#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 = 5381; - for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5) + hostname_hash) ^ hostname[c]; } - - // Distribute all hostname hashes to all processes. - uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); - 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); - cudaCheck(cudaSetDevice(result.local_device_idx)); - ncclUniqueId nccl_id; - if (result.process_rank == 0) { - ncclCheck(ncclGetUniqueId(&nccl_id)); - } - 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)); - return result; -#else - printf("Multi-GPU support is disabled. Using a single GPU."); - return MultiGpuConfig{ - .process_rank = 0, - .num_processes = 1, - .local_device_idx = 0, - }; -#endif -} - -void multi_gpu_config_free(const MultiGpuConfig* multi_gpu_config) { -#ifdef MULTI_GPU - ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm)); - mpiCheck(MPI_Finalize()); -#endif -} - typedef struct { GPT2Config config; // the weights of the model, and their sizes @@ -1692,18 +1710,11 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { if (model_header[1] != 1) { printf("Bad version in model file"); exit(EXIT_FAILURE); } // read in hyperparameters - int maxT, V, L, NH, C; - model->config.max_seq_len = maxT = model_header[2]; - model->config.vocab_size = V = model_header[3]; - model->config.num_layers = L = model_header[4]; - model->config.num_heads = NH = model_header[5]; - model->config.channels = C = model_header[6]; - printf("[GPT-2]\n"); - printf("max_seq_len: %d\n", maxT); - printf("vocab_size: %d\n", V); - printf("num_layers: %d\n", L); - printf("num_heads: %d\n", NH); - printf("channels: %d\n", C); + model->config.max_seq_len = model_header[2]; + model->config.vocab_size = model_header[3]; + model->config.num_layers = model_header[4]; + model->config.num_heads = model_header[5]; + model->config.channels = model_header[6]; // allocate space for all the parameters and read them in fill_in_parameter_sizes(model->param_elements, model->param_sizeof, model->config); @@ -1715,11 +1726,9 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { model->num_parameters_bytes += model->param_elements[i] * model->param_sizeof[i]; } size_t input_model_bytes = model->num_parameters * sizeof(float); - printf("num_parameters: %zu ==> bytes: %zu\n", model->num_parameters, model->num_parameters_bytes); // create memory for model parameters on the device model->params_memory = malloc_and_point_parameters(&model->params, model->param_elements, model->param_sizeof, 1); - printf("allocated %d MiB for model parameters\n", (int)round(model->num_parameters_bytes / (1024 * 1024))); // read in all the parameters from file and copy them to device float* params_memory_cpu = (float*)mallocCheck(input_model_bytes); @@ -1802,7 +1811,7 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, int B, int T) { } model->num_activations = num_activations; model->acts_memory = malloc_and_point_activations(&model->acts, model->act_sizes); - printf("allocated %d MiB for activations\n", (int)round(num_activations * sizeof(floatX) / (1024 * 1024))); + printf0("allocated %d MiB for activations\n", (int)round(num_activations * sizeof(floatX) / (1024 * 1024))); // also create memory for caching inputs and targets cudaCheck(cudaMalloc((void**)&model->inputs, B * T * sizeof(int))); cudaCheck(cudaMalloc((void**)&model->targets, B * T * sizeof(int))); @@ -1918,7 +1927,7 @@ void gpt2_backward(GPT2 *model) { if (model->grads_memory == NULL) { // allocate buffers for weight gradients model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_elements, model->param_sizeof, 1); - printf("allocated %d MiB for parameter gradients\n", (int)round(model->num_parameters * sizeof(floatX) / (1024 * 1024))); + printf0("allocated %d MiB for parameter gradients\n", (int)round(model->num_parameters * sizeof(floatX) / (1024 * 1024))); // we're going to be clever for the activations backward pass. we don't need to exactly // mirror the forward pass acrtivations and we will save memory. size_t bw_act_sizes[NUM_ACTIVATION_TENSORS]; @@ -1931,7 +1940,7 @@ void gpt2_backward(GPT2 *model) { for (size_t i = 0; i < NUM_BACKWARD_TENSORS; i++) { model->num_grad_acts += bw_act_sizes[i]; } - printf("allocated %d MiB for activation gradients\n", (int)round(model->num_grad_acts * sizeof(floatX) / (1024 * 1024))); + printf0("allocated %d MiB for activation gradients\n", (int)round(model->num_grad_acts * sizeof(floatX) / (1024 * 1024))); // init gradients of parameters and activations to zero gpt2_zero_grad(model); } @@ -2074,8 +2083,8 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo cudaCheck(cudaMalloc((void**)&model->v_memory, model->num_parameters * sizeof(float))); cudaCheck(cudaMemset(model->m_memory, 0, model->num_parameters * sizeof(float))); cudaCheck(cudaMemset(model->v_memory, 0, model->num_parameters * sizeof(float))); - printf("allocated %d MiB for AdamW optimizer state m\n", (int)round(model->num_parameters * sizeof(float) / (1024 * 1024))); - printf("allocated %d MiB for AdamW optimizer state v\n", (int)round(model->num_parameters * sizeof(float) / (1024 * 1024))); + printf0("allocated %d MiB for AdamW optimizer state m\n", (int)round(model->num_parameters * sizeof(float) / (1024 * 1024))); + printf0("allocated %d MiB for AdamW optimizer state v\n", (int)round(model->num_parameters * sizeof(float) / (1024 * 1024))); } int block_size = 512; @@ -2366,7 +2375,7 @@ void error_usage() { // ---------------------------------------------------------------------------- // main training loop int main(int argc, char *argv[]) { - MultiGpuConfig multi_gpu_config = multi_gpu_config_init(&argc, &argv); + multi_gpu_config = multi_gpu_config_init(&argc, &argv); // read in the (optional) command line arguments const char* input_dataset_prefix = "data/tiny_shakespeare"; // or e.g. data/TinyStories @@ -2394,19 +2403,19 @@ int main(int argc, char *argv[]) { else if (argv[i][1] == 'g') { genT = atoi(argv[i+1]); } else { error_usage(); } } - printf("+-----------------------+----------------------------------------------------+\n"); - printf("| Parameter | Value |\n"); - printf("+-----------------------+----------------------------------------------------+\n"); - printf("| input dataset prefix | %-50s |\n", input_dataset_prefix); - printf("| output log file | %-50s |\n", output_log_file == NULL ? "NULL" : output_log_file); - printf("| batch size B | %-50d |\n", B); - printf("| sequence length T | %-50d |\n", T); - printf("| learning rate | %-50f |\n", learning_rate); - printf("| val_loss_every | %-50d |\n", val_loss_every); - printf("| val_max_batches | %-50d |\n", val_max_batches); - printf("| sample_every | %-50d |\n", sample_every); - printf("| genT | %-50d |\n", genT); - printf("+-----------------------+----------------------------------------------------+\n"); + printf0("+-----------------------+----------------------------------------------------+\n"); + printf0("| Parameter | Value |\n"); + printf0("+-----------------------+----------------------------------------------------+\n"); + printf0("| input dataset prefix | %-50s |\n", input_dataset_prefix); + printf0("| output log file | %-50s |\n", output_log_file == NULL ? "NULL" : output_log_file); + printf0("| batch size B | %-50d |\n", B); + printf0("| sequence length T | %-50d |\n", T); + printf0("| learning rate | %-50f |\n", learning_rate); + printf0("| val_loss_every | %-50d |\n", val_loss_every); + printf0("| val_max_batches | %-50d |\n", val_max_batches); + printf0("| sample_every | %-50d |\n", sample_every); + printf0("| genT | %-50d |\n", genT); + printf0("+-----------------------+----------------------------------------------------+\n"); // set up the device cudaCheck(cudaSetDevice(multi_gpu_config.local_device_idx)); @@ -2421,20 +2430,20 @@ int main(int argc, char *argv[]) { cublasMath_t cublas_math_mode = enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH; cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode)); cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); - printf("| device | %-50s |\n", deviceProp.name); - printf("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); - printf("+-----------------------+----------------------------------------------------+\n"); + printf0("| device | %-50s |\n", deviceProp.name); + printf0("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); + printf0("+-----------------------+----------------------------------------------------+\n"); // build the GPT-2 model from a checkpoint GPT2 model; gpt2_build_from_checkpoint(&model, "gpt2_124M.bin"); - printf("| max_sequence_length T | %-50d |\n", model.config.max_seq_len); - printf("| vocab_size V | %-50d |\n", model.config.vocab_size); - printf("| num_layers L | %-50d |\n", model.config.num_layers); - printf("| num_heads NH | %-50d |\n", model.config.num_heads); - printf("| channels C | %-50d |\n", model.config.channels); - printf("| num_parameters | %-50zu |\n", model.num_parameters); - printf("+-----------------------+----------------------------------------------------+\n"); + printf0("| max_sequence_length T | %-50d |\n", model.config.max_seq_len); + printf0("| vocab_size V | %-50d |\n", model.config.vocab_size); + printf0("| num_layers L | %-50d |\n", model.config.num_layers); + printf0("| num_heads NH | %-50d |\n", model.config.num_heads); + printf0("| channels C | %-50d |\n", model.config.channels); + printf0("| num_parameters | %-50zu |\n", model.num_parameters); + printf0("+-----------------------+----------------------------------------------------+\n"); // build DataLoaders for both train and val char train_tokens_filename[128]; @@ -2448,12 +2457,17 @@ int main(int argc, char *argv[]) { dataloader_init(&val_loader, &multi_gpu_config, val_tokens_filename, B, T); int train_num_batches = train_loader.num_batches; // let's do 1 epoch by default for now int val_num_batches = train_loader.num_batches < val_max_batches ? train_loader.num_batches : val_max_batches; - printf("| train_num_batches | %-50d |\n", train_num_batches); - printf("| val_num_batches | %-50d |\n", val_num_batches); - printf("+-----------------------+----------------------------------------------------+\n"); + printf0("| train_num_batches | %-50d |\n", train_num_batches); + printf0("| val_num_batches | %-50d |\n", val_num_batches); + printf0("+-----------------------+----------------------------------------------------+\n"); - // print model parameter allocations from gpt2_build_from_checkpoint down here to not mess up our table above - printf("allocated %d MiB for model parameters\n", (int)round(model.num_parameters_bytes / (1024 * 1024))); + // pretty print in a table the multi-gpu configuration as well + printf0("| num_processes | %-50d |\n", multi_gpu_config.num_processes); + printf0("+-----------------------+----------------------------------------------------+\n"); + + // more prints related to allocations from gpt2_build_from_checkpoint down here to not mess up our table above + printf0("num_parameters: %zu ==> bytes: %zu\n", model.num_parameters, model.num_parameters_bytes); + printf0("allocated %d MiB for model parameters\n", (int)round(model.num_parameters_bytes / (1024 * 1024))); // set up the Logger Logger logger; @@ -2491,7 +2505,7 @@ int main(int argc, char *argv[]) { } // once in a while do model inference to print generated text - if (multi_gpu_config.process_rank == 0 && step > 0 && step % sample_every == 0 || last_step) { + if (multi_gpu_config.process_rank == 0 && (step > 0 && (step % sample_every) == 0 || last_step)) { // fill up gen_tokens with the GPT2_EOT, which kicks off the generation for(int i = 0; i < B * T; ++i) { gen_tokens[i] = GPT2_EOT; @@ -2550,12 +2564,12 @@ int main(int argc, char *argv[]) { clock_gettime(CLOCK_MONOTONIC, &end); double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9; total_sum_iteration_time_s += time_elapsed_s; - int tokens_per_second = (B * T) / time_elapsed_s; - printf("step %4d/%d: train loss %f (acc %f) (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, model.accumulated_mean_loss, time_elapsed_s * 1000, tokens_per_second); + int tokens_per_second = multi_gpu_config.num_processes * (B * T) / time_elapsed_s; + printf0("step %4d/%d: train loss %f (acc %f) (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, model.accumulated_mean_loss, time_elapsed_s * 1000, tokens_per_second); logger_log_train(&logger, step, model.mean_loss); } // add a total average, for optimizations that are only mild improvements - printf("total average iteration time: %f ms\n", total_sum_iteration_time_s / train_num_batches * 1000); + printf0("total average iteration time: %f ms\n", total_sum_iteration_time_s / train_num_batches * 1000); // free dataloader_free(&train_loader); diff --git a/train_gpt2_fp32.cu b/train_gpt2_fp32.cu index 7eda056..a0b5b1b 100644 --- a/train_gpt2_fp32.cu +++ b/train_gpt2_fp32.cu @@ -1331,6 +1331,7 @@ int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { // 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.