From d49e8301ebe330c2b25c2535fcaafda581fdc851 Mon Sep 17 00:00:00 2001 From: Petr Zhizhin Date: Mon, 22 Apr 2024 22:12:33 +0200 Subject: [PATCH] [Multi-GPU] llm.c now runs on multiple GPUs with NCCL --- Makefile | 27 ++++- README.md | 12 +++ dev/cuda/nccl_all_reduce.cu | 196 +++++++++++++++++++++++++++++++++++ train_gpt2.cu | 200 +++++++++++++++++++++++++++++++++--- train_gpt2_fp32.cu | 180 +++++++++++++++++++++++++++++--- 5 files changed, 579 insertions(+), 36 deletions(-) create mode 100644 dev/cuda/nccl_all_reduce.cu diff --git a/Makefile b/Makefile index 7ef5bdf..5e6ecbc 100644 --- a/Makefile +++ b/Makefile @@ -11,6 +11,8 @@ NVCC := $(shell which nvcc 2>/dev/null) # NVCC flags NVCC_FLAGS = -O3 --use_fast_math NVCC_LDFLAGS = -lcublas -lcublasLt +NCLL_INCLUDES = +NVCC_LDLIBS = # Function to test if the compiler accepts a given flag. define check_and_add_flag @@ -65,6 +67,23 @@ else endif endif +ifeq ($(NO_MULTI_GPU), 1) + $(info Multi-GPU (OpenMPI + NCCL) is manually disabled) +else + # Detect if running on macOS or Linux + 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/ ] && echo "exists"), exists) + $(info Adding OpenMPI 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 + NVCC_FLAGS += -DMULTI_GPU + else + $(warning OpenMPI is not found, disabling multi-GPU support) + endif +endif + # PHONY means these targets will always be executed .PHONY: all train_gpt2 test_gpt2 train_gpt2cu test_gpt2cu train_gpt2fp32cu test_gpt2fp32cu @@ -88,16 +107,16 @@ test_gpt2: test_gpt2.c $(CC) $(CFLAGS) $(INCLUDES) $(LDFLAGS) $< $(LDLIBS) -o $@ train_gpt2cu: train_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) -o $@ + $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) -o $@ train_gpt2fp32cu: train_gpt2_fp32.cu - $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) -o $@ + $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) -o $@ test_gpt2cu: test_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) -o $@ + $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) -o $@ test_gpt2fp32cu: test_gpt2_fp32.cu - $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) -o $@ + $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) -o $@ profile_gpt2cu: profile_gpt2.cu $(NVCC) $(NVCC_FLAGS) -lineinfo $< $(NVCC_LDFLAGS) -o $@ diff --git a/README.md b/README.md index 3cbc405..609c00c 100644 --- a/README.md +++ b/README.md @@ -34,6 +34,18 @@ OMP_NUM_THREADS=8 ./train_gpt2 The above lines (1) download the [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset, tokenize it with the GPT-2 Tokenizer, (2) download and save the GPT-2 (124M) weights, (3) init from them in C and train for 40 steps on tineshakespeare with AdamW (using batch size 4, context length only 64), evaluate validation loss, and sample some text. Honestly, unless you have a beefy CPU (and can crank up the number of OMP threads in the launch command), you're not going to get that far on CPU training LLMs, but it might be a good demo/reference. +## quick start (multiple GPUs) + +The "Maybe I'm GPU poor, but at least I have a couple". No worries, run: +``` +sudo apt install openmpi-bin openmpi-doc libopenmpi-dev +pip install -r requirements.txt +python prepro_tinyshakespeare.py +python train_gpt2.py +make train_gpt2fp32cu +mpirun -np ./train_gpt2fp32cu +``` + ## training: more detail Download and tokenize a dataset. The [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset is the fastest to download and tokenize: diff --git a/dev/cuda/nccl_all_reduce.cu b/dev/cuda/nccl_all_reduce.cu new file mode 100644 index 0000000..7559c5b --- /dev/null +++ b/dev/cuda/nccl_all_reduce.cu @@ -0,0 +1,196 @@ +/* + +A simple test of NCCL capabilities. +Fills a vector with 1s on the first GPU, 2s on the second, etc. +Then aggregates the values in the resulting vectors. + +Compile example: +nvcc -lmpi -lnccl -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux-gnu/openmpi/lib/ nccl_all_reduce.cu -o nccl_all_reduce + +Run on 2 local GPUs (set -np to a different value to change GPU count): +mpirun -np 2 ./nccl_all_reduce + +*/ + +#include "common.h" +#include +#include +#include +#include +#include +#include +#include +#include + +void nccl_check(ncclResult_t status, const char *file, int line) { + if (status != ncclSuccess) { + printf("[NCCL ERROR] at file %s:%d:\n%s\n", file, line, + ncclGetErrorString(status)); + exit(EXIT_FAILURE); + } +} +#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__)) + +// Sets a vector to a predefined value +__global__ void set_vector(float *data, int N, float value) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + + // Check for out-of-bounds access + if (i < N) { + data[i] = value; + } +} + +size_t cdiv(size_t a, size_t b) { return (a + b - 1) / b; } + +// Parameters specific to training on multiple GPUs. +typedef struct { + int process_rank; // Rank of this process among all MPI processes on all hosts. 0 if no multi-GPU. + int num_processes; // Total number of processes on all hosts. 1 if no multi-GPU. + int local_device_idx; // This process GPU index on current machine. 0 if no multi-GPU. + ncclComm_t nccl_comm; // NCCL communication primitive, used for collective mutli-GPU work. +} MultiGpuConfig; + +// 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; +} + +MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) { + // 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); + printf("[Process rank %d] Using GPU %d\n", result.process_rank, result.local_device_idx); + 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; +} + +void multi_gpu_config_free(const MultiGpuConfig* multi_gpu_config) { + ncclCommDestroy(multi_gpu_config->nccl_comm); + mpiCheck(MPI_Finalize()); +} + +float get_mean(float *arr, size_t size, int process_rank) { + double sum = 0.0; + for (size_t i = 0; i < size; ++i) { + sum += arr[i]; + } + return sum / size; +} + +int main(int argc, char **argv) { + // Some constants + const size_t all_reduce_buffer_size = 32 * 1024 * 1024; + const size_t threads_per_block = 1024; + + MultiGpuConfig multi_gpu_config = multi_gpu_config_init(&argc, &argv); + + // Allocating buffers on each of the devices. + float *all_reduce_buffer; + cudaCheck( + cudaMalloc(&all_reduce_buffer, all_reduce_buffer_size * sizeof(float))); + + int n_blocks = cdiv(all_reduce_buffer_size, threads_per_block); + // Set the allocated memory to a defined value. + set_vector<<>>( + all_reduce_buffer, all_reduce_buffer_size, + (float)(multi_gpu_config.process_rank + 1)); + cudaCheck(cudaGetLastError()); + + float *all_reduce_buffer_host = + (float *)malloc(all_reduce_buffer_size * sizeof(float)); + + cudaCheck(cudaMemcpy(all_reduce_buffer_host, all_reduce_buffer, + sizeof(float) * all_reduce_buffer_size, + cudaMemcpyDeviceToHost)); + + printf("[Process rank %d] average value before all reduce is %.6f\n", multi_gpu_config.process_rank, + get_mean(all_reduce_buffer_host, all_reduce_buffer_size, + multi_gpu_config.process_rank)); + + float *all_reduce_buffer_recv; + cudaCheck(cudaMalloc(&all_reduce_buffer_recv, + all_reduce_buffer_size * sizeof(float))); + + ncclCheck(ncclAllReduce( + (const void *)all_reduce_buffer, (void *)all_reduce_buffer_recv, + all_reduce_buffer_size, ncclFloat, ncclSum, multi_gpu_config.nccl_comm, 0)); + + + cudaCheck(cudaMemcpy(all_reduce_buffer_host, all_reduce_buffer_recv, + sizeof(float) * all_reduce_buffer_size, + cudaMemcpyDeviceToHost)); + + float all_reduce_mean_value = get_mean(all_reduce_buffer_host, all_reduce_buffer_size, multi_gpu_config.process_rank); + + printf("[Process rank %d] average value after all reduce is %.6f\n", multi_gpu_config.process_rank, all_reduce_mean_value); + + float expected_all_reduce_mean_value = 0.0; + for (int i = 0; i != multi_gpu_config.num_processes; ++i) { + expected_all_reduce_mean_value += i + 1; + } + if (abs(expected_all_reduce_mean_value - all_reduce_mean_value) > 1e-5) { + printf("[Process rank %d] ERROR: Unexpected all reduce value: %.8f, expected %.8f\n", multi_gpu_config.process_rank, all_reduce_mean_value, expected_all_reduce_mean_value); + } else { + printf("[Process rank %d] Checked against expected mean value. All good!\n", multi_gpu_config.process_rank); + } + + free(all_reduce_buffer_host); + cudaCheck(cudaFree(all_reduce_buffer)); + multi_gpu_config_free(&multi_gpu_config); +} diff --git a/train_gpt2.cu b/train_gpt2.cu index 913a653..67732ed 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -32,6 +32,11 @@ and fp8 (coming soon^TM). #include #include +#ifdef MULTI_GPU +#include +#include +#endif + // ---------------------------------------------------------------------------- // CUDA precision settings @@ -45,6 +50,11 @@ typedef float floatN; #define CUBLAS_LOWP CUDA_R_16BF #define CUBLAS_LOWP_COMPUTE CUBLAS_COMPUTE_32F +#ifdef MULTI_GPU +const ncclDataType_t ncclFloatX = ncclBfloat16; +const ncclDataType_t ncclFloatN = ncclFloat; +#endif + // use fp16 (note: this may require gradient scaler, currently not implemented!) #elif defined(ENABLE_FP16) typedef half floatX; @@ -52,12 +62,23 @@ typedef float floatN; #define CUBLAS_LOWP CUDA_R_16F #define CUBLAS_LOWP_COMPUTE CUBLAS_COMPUTE_32F +#ifdef MULTI_GPU +const ncclDataType_t ncclFloatX = ncclHalf; +const ncclDataType_t ncclFloatN = ncclFloat; +#endif + // fallback for fp32 #else typedef float floatX; typedef float floatN; #define CUBLAS_LOWP CUDA_R_32F #define CUBLAS_LOWP_COMPUTE cublas_compute_type // auto-select FP32 vs TF32 + +#ifdef MULTI_GPU +const ncclDataType_t ncclFloatX = ncclFloat; +const ncclDataType_t ncclFloatN = ncclFloat; +#endif + #endif // ---------------------------------------------------------------------------- @@ -95,6 +116,27 @@ void cublasCheck(cublasStatus_t status, const char *file, int line) } #define cublasCheck(status) { cublasCheck((status), __FILE__, __LINE__); } +#ifdef MULTI_GPU +void nccl_check(ncclResult_t status, const char *file, int line) { + if (status != ncclSuccess) { + printf("[NCCL ERROR] at file %s:%d:\n%s\n", file, line, ncclGetErrorString(status)); + exit(EXIT_FAILURE); + } +} +#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 + // GPU helper functions for atomicAdd on smaller than 32-bit types __device__ void atomicAddX(__nv_bfloat16* addr, __nv_bfloat16 val) { uintptr_t ptr_val = reinterpret_cast(addr); @@ -1516,6 +1558,89 @@ 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 @@ -1546,6 +1671,7 @@ typedef struct { int* inputs; // the input tokens for the current forward pass int* targets; // the target tokens for the current forward pass float mean_loss; // after a forward pass with targets, will be populated with the mean loss + float accumulated_mean_loss; // Mean loss after aggregating it on all GPUs floatX* cpu_losses; // CPU buffer to copy the losses to, allocated with cudaMallocHost unsigned long long rng_state; // the RNG state for seeding stochastic rounding etc. } GPT2; @@ -1898,6 +2024,41 @@ void gpt2_backward(GPT2 *model) { encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C); } +// Compute a mean of a single CPU value across all GPU processes. No-op when multi-GPU is disabled. +float multi_gpu_cpu_float_mean(float value, const MultiGpuConfig* multi_gpu_config) { +#ifdef MULTI_GPU + // MPI doesn't support all reduce with mean, so we sum up, then divide. + float result; + mpiCheck(MPI_Allreduce(&value, &result, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD)); + return result / multi_gpu_config->num_processes; +#else + return value; +#endif +} + +// Averages out the loss and gradients across all GPUs. No-op when multi-GPU is disabled. +void gpt2_mutli_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) { + // Average all losses. + model->accumulated_mean_loss = multi_gpu_cpu_float_mean(model->mean_loss, multi_gpu_config); +#ifdef MULTI_GPU + // Average all gradients. + char* grads_memory_iterator = (char*)model->grads_memory; + for (int i = 0; i < NUM_PARAMETER_TENSORS; ++i) { + int current_param_sizeof = model->param_sizeof[i]; + int current_param_elements = model->param_elements[i]; + ncclDataType_t data_type = current_param_sizeof == sizeof(floatX) ? ncclFloatX : ncclFloatN; + ncclCheck(ncclAllReduce(grads_memory_iterator, grads_memory_iterator, + current_param_elements, + data_type, ncclAvg, + multi_gpu_config->nccl_comm, + // use 0 for default stream (all other computations use this stream) + /*stream=*/0)); + grads_memory_iterator += current_param_elements * current_param_sizeof; + } + assert(grads_memory_iterator == (char*)model->grads_memory + model->num_parameters_bytes); +#endif +} + void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) { // reference: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html @@ -1970,6 +2131,10 @@ void gpt2_free(GPT2 *model) { // data loader lite: returns random batches of data from a file of integers typedef struct { + // Distributed data parallel specifics. + // Each worker loads it's own chunk of data. + int process_rank; + int num_processes; // hyperparameters int B; int T; @@ -1985,7 +2150,9 @@ typedef struct { int num_batches; } DataLoader; -void dataloader_init(DataLoader *loader, const char* filename, int B, int T) { +void dataloader_init(DataLoader *loader, const MultiGpuConfig* multi_gpu_config, const char* filename, int B, int T) { + loader->process_rank = multi_gpu_config->process_rank; + loader->num_processes = multi_gpu_config->num_processes; loader->B = B; loader->T = T; @@ -2000,7 +2167,7 @@ void dataloader_init(DataLoader *loader, const char* filename, int B, int T) { printf("Error: file size is too small for the batch size and sequence length\n"); exit(EXIT_FAILURE); } - loader->current_position = 0; // start at the beginning + loader->current_position = loader->process_rank * B * T * sizeof(int); // start at the beginning // allocate space for B*T + 1 integers to store the inputs and targets // Using CUDA CPU pinned memory for faster PCI Express transfers to GPU @@ -2008,7 +2175,7 @@ void dataloader_init(DataLoader *loader, const char* filename, int B, int T) { cudaMallocHost((void**)&loader->batch, (B * T + 1) * sizeof(int)); loader->inputs = loader->batch; loader->targets = loader->batch + 1; // targets are shifted by one - loader->num_batches = loader->file_size / (B * T * sizeof(int)); + loader->num_batches = loader->file_size / (loader->num_processes * B * T * sizeof(int)); } void dataloader_reset(DataLoader *loader) { @@ -2019,14 +2186,14 @@ void dataloader_next_batch(DataLoader *loader) { int B = loader->B; int T = loader->T; // if we are at the end of the file, loop back to the beginning - if (loader->current_position + (B*T+1) * sizeof(int) > loader->file_size) { - loader->current_position = 0; + if (loader->current_position + (loader->num_processes * B * T + 1) * sizeof(int) > loader->file_size) { + loader->current_position = loader->process_rank * B * T * sizeof(int); } // read the B*T+1 integers from the file into batch fseek(loader->tokens_file, loader->current_position, SEEK_SET); freadCheck(loader->batch, sizeof(int), B*T+1, loader->tokens_file); - // advance the current position by B*T integers - loader->current_position += B*T * sizeof(int); + // advance the current position by B*T*num_processes integers + loader->current_position += loader->num_processes * B * T * sizeof(int); } void dataloader_free(DataLoader *loader) { @@ -2193,6 +2360,7 @@ void error_usage() { // ---------------------------------------------------------------------------- // main training loop int main(int argc, char *argv[]) { + MultiGpuConfig 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 @@ -2211,7 +2379,7 @@ int main(int argc, char *argv[]) { // read in the args if (argv[i][1] == 'i') { input_dataset_prefix = argv[i+1]; } else if (argv[i][1] == 'o') { output_log_file = argv[i+1]; } - else if (argv[i][1] == 'b') { B = atoi(argv[i+1]); } + else if (argv[i][1] == 'b') { B = atoi(argv[i+1]); } // Per-GPU batch size else if (argv[i][1] == 't') { T = atoi(argv[i+1]); } else if (argv[i][1] == 'l') { learning_rate = atof(argv[i+1]); } else if (argv[i][1] == 'v') { val_loss_every = atoi(argv[i+1]); } @@ -2235,10 +2403,9 @@ int main(int argc, char *argv[]) { printf("+-----------------------+----------------------------------------------------+\n"); // set up the device - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); + cudaCheck(cudaSetDevice(multi_gpu_config.local_device_idx)); cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, deviceIdx); + cudaGetDeviceProperties(&deviceProp, multi_gpu_config.local_device_idx); // setup cuBLAS and cuBLASLt cublasCheck(cublasCreate(&cublas_handle)); cublasCheck(cublasLtCreate(&cublaslt_handle)); @@ -2270,9 +2437,9 @@ int main(int argc, char *argv[]) { sprintf(train_tokens_filename, "%s_train.bin", input_dataset_prefix); sprintf(val_tokens_filename, "%s_val.bin", input_dataset_prefix); DataLoader train_loader; - dataloader_init(&train_loader, train_tokens_filename, B, T); + dataloader_init(&train_loader, &multi_gpu_config, train_tokens_filename, B, T); DataLoader val_loader; - dataloader_init(&val_loader, val_tokens_filename, B, T); + 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); @@ -2312,12 +2479,13 @@ int main(int argc, char *argv[]) { val_loss += model.mean_loss; } val_loss /= val_num_batches; + val_loss = multi_gpu_cpu_float_mean(val_loss, &multi_gpu_config); printf("val loss %f\n", val_loss); logger_log_val(&logger, step, val_loss); } // once in a while do model inference to print generated text - if (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; @@ -2370,13 +2538,14 @@ int main(int argc, char *argv[]) { gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T); gpt2_zero_grad(&model); gpt2_backward(&model); + gpt2_mutli_gpu_accumulate(&model, &multi_gpu_config); gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1); cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings 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 (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, time_elapsed_s * 1000, tokens_per_second); + 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); logger_log_train(&logger, step, model.mean_loss); } // add a total average, for optimizations that are only mild improvements @@ -2394,6 +2563,7 @@ int main(int argc, char *argv[]) { cublasCheck(cublasDestroy(cublas_handle)); cublasCheck(cublasLtDestroy(cublaslt_handle)); logger_free(&logger); + multi_gpu_config_free(&multi_gpu_config); return 0; } diff --git a/train_gpt2_fp32.cu b/train_gpt2_fp32.cu index c2027e8..7eda056 100644 --- a/train_gpt2_fp32.cu +++ b/train_gpt2_fp32.cu @@ -27,6 +27,11 @@ the layernorms are connected to the residuals so we += in layernorm backward. #include #include +#ifdef MULTI_GPU +#include +#include +#endif + // ---------------------------------------------------------------------------- // CUDA utils @@ -53,6 +58,27 @@ void cublasCheck(cublasStatus_t status, const char *file, int line) } #define cublasCheck(status) { cublasCheck((status), __FILE__, __LINE__); } +#ifdef MULTI_GPU +void nccl_check(ncclResult_t status, const char *file, int line) { + if (status != ncclSuccess) { + printf("[NCCL ERROR] at file %s:%d:\n%s\n", file, line, ncclGetErrorString(status)); + exit(EXIT_FAILURE); + } +} +#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 + // cuBLAS workspace. Hardcoding to 32MiB but only Hopper needs 32, for others 4 is OK static size_t cublaslt_workspace_size = 32 * 1024 * 1024; static void* cublaslt_workspace = NULL; @@ -1275,6 +1301,89 @@ float* 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 on all hosts. 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 @@ -1304,6 +1413,7 @@ typedef struct { int* targets; // the target tokens for the current forward pass float mean_loss; // after a forward pass with targets, will be populated with the mean loss float* cpu_losses; // CPU buffer to copy the losses to, allocated with cudaMallocHost + float accumulated_mean_loss; // Mean loss after aggregating it on all GPUs } GPT2; void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { @@ -1621,6 +1731,33 @@ void gpt2_backward(GPT2 *model) { encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C); } +// Compute a mean of a single CPU value across all GPU processes. No-op when multi-GPU is disabled. +float multi_gpu_cpu_float_mean(float value, const MultiGpuConfig* multi_gpu_config) { +#ifdef MULTI_GPU + // MPI doesn't support all reduce with mean, so we sum up, then divide. + float result; + mpiCheck(MPI_Allreduce(&value, &result, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD)); + return result / multi_gpu_config->num_processes; +#else + return value; +#endif +} + +// Averages out the loss and gradients across all GPUs. No-op when multi-GPU is disabled. +void gpt2_mutli_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) { + // Average all losses. + model->accumulated_mean_loss = multi_gpu_cpu_float_mean(model->mean_loss, multi_gpu_config); +#ifdef MULTI_GPU + // Average all gradients. + ncclCheck(ncclAllReduce(model->grads_memory, model->grads_memory, + model->num_parameters, + ncclFloat, ncclAvg, + multi_gpu_config->nccl_comm, + // use 0 for default stream (all other computations use this stream) + /*stream=*/0)); +#endif +} + void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) { // reference: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html @@ -1663,6 +1800,10 @@ void gpt2_free(GPT2 *model) { // data loader lite: returns random batches of data from a file of integers typedef struct { + // Distributed data parallel specifics. + // Each worker loads it's own chunk of data. + int process_rank; + int num_processes; // hyperparameters int B; int T; @@ -1678,7 +1819,9 @@ typedef struct { long num_batches; } DataLoader; -void dataloader_init(DataLoader *loader, const char* filename, int B, int T) { +void dataloader_init(DataLoader *loader, const MultiGpuConfig* multi_gpu_config, const char* filename, int B, int T) { + loader->process_rank = multi_gpu_config->process_rank; + loader->num_processes = multi_gpu_config->num_processes; loader->B = B; loader->T = T; @@ -1693,7 +1836,7 @@ void dataloader_init(DataLoader *loader, const char* filename, int B, int T) { printf("Error: file size is too small for the batch size and sequence length\n"); exit(EXIT_FAILURE); } - loader->current_position = 0; // start at the beginning + loader->current_position = loader->process_rank * B * T * sizeof(int); // start at the beginning // allocate space for B*T + 1 integers to store the inputs and targets // Using CUDA CPU pinned memory for faster PCI Express transfers to GPU @@ -1701,25 +1844,24 @@ void dataloader_init(DataLoader *loader, const char* filename, int B, int T) { cudaMallocHost((void**)&loader->batch, (B * T + 1) * sizeof(int)); loader->inputs = loader->batch; loader->targets = loader->batch + 1; // targets are shifted by one - loader->num_batches = loader->file_size / (B * T * sizeof(int)); + loader->num_batches = loader->file_size / (loader->num_processes * B * T * sizeof(int)); } void dataloader_reset(DataLoader *loader) { - loader->current_position = 0; + loader->current_position = loader->process_rank * loader->B * loader->T * sizeof(int); } void dataloader_next_batch(DataLoader *loader) { int B = loader->B; int T = loader->T; // if we are at the end of the file, loop back to the beginning - if (loader->current_position + (B*T+1) * sizeof(int) > loader->file_size) { - loader->current_position = 0; + if (loader->current_position + (loader->num_processes * B * T + 1) * sizeof(int) > loader->file_size) { + loader->current_position = loader->process_rank * B * T * sizeof(int); } // read the B*T+1 integers from the file into batch fseek(loader->tokens_file, loader->current_position, SEEK_SET); freadCheck(loader->batch, sizeof(int), B*T+1, loader->tokens_file); - // advance the current position by B*T integers - loader->current_position += B*T * sizeof(int); + loader->current_position += loader->num_processes * B * T * sizeof(int); } void dataloader_free(DataLoader *loader) { @@ -1897,6 +2039,7 @@ void error_usage() { // ---------------------------------------------------------------------------- // main training loop int main(int argc, char *argv[]) { + MultiGpuConfig 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 @@ -1915,7 +2058,7 @@ int main(int argc, char *argv[]) { // read in the args if (argv[i][1] == 'i') { input_dataset_prefix = argv[i+1]; } else if (argv[i][1] == 'o') { output_log_file = argv[i+1]; } - else if (argv[i][1] == 'b') { B = atoi(argv[i+1]); } + else if (argv[i][1] == 'b') { B = atoi(argv[i+1]); } // Per-GPU batch size else if (argv[i][1] == 't') { T = atoi(argv[i+1]); } else if (argv[i][1] == 'l') { learning_rate = atof(argv[i+1]); } else if (argv[i][1] == 'v') { val_loss_every = atoi(argv[i+1]); } @@ -1938,11 +2081,8 @@ int main(int argc, char *argv[]) { printf("| genT | %-50d |\n", genT); printf("+-----------------------+----------------------------------------------------+\n"); - // set up the device - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, deviceIdx); + cudaGetDeviceProperties(&deviceProp, multi_gpu_config.local_device_idx); // setup cuBLAS and cuBLASLt cublasCheck(cublasCreate(&cublas_handle)); cublasCheck(cublasLtCreate(&cublaslt_handle)); @@ -1952,6 +2092,9 @@ 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("| process rank | %-50d |\n", multi_gpu_config.process_rank); + printf("| number of processes | %-50d |\n", multi_gpu_config.num_processes); + printf("| local device index | %-50d |\n", multi_gpu_config.local_device_idx); printf("| device | %-50s |\n", deviceProp.name); printf("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); printf("+-----------------------+----------------------------------------------------+\n"); @@ -1974,9 +2117,9 @@ int main(int argc, char *argv[]) { sprintf(train_tokens_filename, "%s_train.bin", input_dataset_prefix); sprintf(val_tokens_filename, "%s_val.bin", input_dataset_prefix); DataLoader train_loader; - dataloader_init(&train_loader, train_tokens_filename, B, T); + dataloader_init(&train_loader, &multi_gpu_config, train_tokens_filename, B, T); DataLoader val_loader; - dataloader_init(&val_loader, val_tokens_filename, B, T); + 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); @@ -2015,12 +2158,13 @@ int main(int argc, char *argv[]) { val_loss += model.mean_loss; } val_loss /= val_num_batches; + val_loss = multi_gpu_cpu_float_mean(val_loss, &multi_gpu_config); printf("val loss %f\n", val_loss); logger_log_val(&logger, step, val_loss); } // once in a while do model inference to print generated text - if (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; @@ -2068,13 +2212,14 @@ int main(int argc, char *argv[]) { gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T); gpt2_zero_grad(&model); gpt2_backward(&model); + gpt2_mutli_gpu_accumulate(&model, &multi_gpu_config); gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1); cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings 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 (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, time_elapsed_s * 1000, tokens_per_second); + 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); logger_log_train(&logger, step, model.mean_loss); } // add a total average, for optimizations that are only mild improvements @@ -2091,6 +2236,7 @@ int main(int argc, char *argv[]) { cublasCheck(cublasDestroy(cublas_handle)); cublasCheck(cublasLtDestroy(cublaslt_handle)); logger_free(&logger); + multi_gpu_config_free(&multi_gpu_config); return 0; }