From 0d52d2a3d797366fdcdebdb7976a8319c453cbd1 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 24 Jul 2024 15:46:30 +0200 Subject: [PATCH 1/7] fall back to cudaMallocManaged for optimizer states if we're out of memory --- llmc/cuda_common.h | 4 ++-- llmc/cuda_utils.cuh | 23 +++++++++++++++++++++++ train_gpt2.cu | 6 +++--- 3 files changed, 28 insertions(+), 5 deletions(-) diff --git a/llmc/cuda_common.h b/llmc/cuda_common.h index 006ad30..2b07aab 100644 --- a/llmc/cuda_common.h +++ b/llmc/cuda_common.h @@ -49,13 +49,13 @@ constexpr std::bool_constant False; // Error checking // CUDA error checking -inline void cudaCheck(cudaError_t error, const char *file, int line) { +inline void cudaCheck_(cudaError_t error, const char *file, int line) { if (error != cudaSuccess) { printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, cudaGetErrorString(error)); exit(EXIT_FAILURE); } }; -#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__)) +#define cudaCheck(err) (cudaCheck_(err, __FILE__, __LINE__)) // like cudaFree, but checks for errors _and_ resets the pointer. template diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 0ce728e..22cdc6b 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -205,6 +205,29 @@ void global_sum_deterministic(float* result, const Float* values, int count, cud cudaCheck(cudaGetLastError()); } +// ---------------------------------------------------------------------------- +// memory management + +// allocate memory, preferrably on the +void cudaMallocConditionallyManaged(void** out, size_t bytes, const char *file, int line) { + size_t free, total; + cudaCheck(cudaMemGetInfo(&free, &total)); + // check if we have enough space to pin the memory to device (with 1% slack) + if(100 * free < 99 * bytes) { + cudaCheck_(cudaMalloc((void**)out, bytes), file, line); + } else { + // if not, fallback to a managed allocation. It will be slower, but at least + // it won't crash. + fprintf(stderr, "[WARN] Not enough space to allocate %zu bytes on device.\n" + " Falling back to managed allocation.\n Speed may be negatively affected.", + bytes); + cudaCheck_(cudaMallocManaged((void**)out, bytes), file, line); + } +} + +#define cudaMallocConditionallyManaged(out, bytes)\ +(cudaMallocConditionallyManaged((void**)out, bytes, __FILE__, __LINE__)) + // ---------------------------------------------------------------------------- // Random Number Generation used in Stochastic Rounding diff --git a/train_gpt2.cu b/train_gpt2.cu index 16f8013..343343e 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -393,13 +393,13 @@ void gpt2_allocate_state(GPT2 *model, int B, int T) { printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); assert(model->m_memory == nullptr); assert(model->v_memory == nullptr); - cudaCheck(cudaMalloc((void**)&model->m_memory, shard_num_parameters * sizeof(float))); - cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float))); + cudaMallocConditionallyManaged((void**)&model->m_memory, shard_num_parameters * sizeof(float)); + cudaMallocConditionallyManaged((void**)&model->v_memory, shard_num_parameters * sizeof(float)); if (model->use_master_weights == 1) { assert(model->master_weights == nullptr); 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))); + cudaMallocConditionallyManaged((void**) &model->master_weights, shard_num_parameters * sizeof(float)); } size_t free, total; From c8457572fab5272d3951140d819598fd59aeff34 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Wed, 24 Jul 2024 16:55:09 +0200 Subject: [PATCH 2/7] just try to allocate on device; fallback if that fails --- llmc/cuda_utils.cuh | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 22cdc6b..81e4066 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -210,18 +210,20 @@ void global_sum_deterministic(float* result, const Float* values, int count, cud // allocate memory, preferrably on the void cudaMallocConditionallyManaged(void** out, size_t bytes, const char *file, int line) { - size_t free, total; - cudaCheck(cudaMemGetInfo(&free, &total)); - // check if we have enough space to pin the memory to device (with 1% slack) - if(100 * free < 99 * bytes) { - cudaCheck_(cudaMalloc((void**)out, bytes), file, line); - } else { - // if not, fallback to a managed allocation. It will be slower, but at least + // try to allocate `bytes` on device + cudaError_t err = cudaMalloc(out, bytes); + if(err == cudaErrorMemoryAllocation) { + // if that fails, fallback to a managed allocation. It will be slower, but at least // it won't crash. - fprintf(stderr, "[WARN] Not enough space to allocate %zu bytes on device.\n" - " Falling back to managed allocation.\n Speed may be negatively affected.", - bytes); - cudaCheck_(cudaMallocManaged((void**)out, bytes), file, line); + fprintf(stderr, "[WARN] Not enough space to allocate %zu MiB on device.\n" + " Falling back to managed allocation.\n" + " Speed may be negatively affected.\n", + bytes / 1024 / 1024); + // reset the error before the next API call + cudaGetLastError(); + cudaCheck_(cudaMallocManaged(out, bytes), file, line); + } else { + cudaCheck_(err, file, line); } } From f72c1f2c6428e77db130a7db859d599ef359beb4 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Thu, 25 Jul 2024 01:16:28 +0200 Subject: [PATCH 3/7] hint to host --- llmc/cuda_utils.cuh | 1 + 1 file changed, 1 insertion(+) diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 81e4066..d3df74d 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -222,6 +222,7 @@ void cudaMallocConditionallyManaged(void** out, size_t bytes, const char *file, // reset the error before the next API call cudaGetLastError(); cudaCheck_(cudaMallocManaged(out, bytes), file, line); + cudaCheck_(cudaMemAdvise(*out, bytes, cudaMemAdviseSetPreferredLocation, cudaCpuDeviceId), file, line); } else { cudaCheck_(err, file, line); } From 2882ec6b9e992cc28604bdfa327c6f82c0e19aea Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Aug 2024 17:22:30 +0000 Subject: [PATCH 4/7] fix makefile for multigpu setups --- Makefile | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 6fa511d..73b8372 100644 --- a/Makefile +++ b/Makefile @@ -49,7 +49,9 @@ endif ifneq ($(CI),true) # if not in CI, then use the GPU query ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY= ifneq ($(call file_exists_in_path, nvidia-smi),) - GPU_COMPUTE_CAPABILITY = $(shell nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed 's/\.//g') + # Get the compute capabilities of all GPUs + # Remove decimal points, sort numerically in ascending order, and select the first (lowest) value + GPU_COMPUTE_CAPABILITY=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed 's/\.//g' | sort -n | head -n 1) GPU_COMPUTE_CAPABILITY := $(strip $(GPU_COMPUTE_CAPABILITY)) endif endif From e6856bc5664ca4b43d6c87883a0c5cadfc202cf6 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Aug 2024 17:23:42 +0000 Subject: [PATCH 5/7] fallback to memory allocation of m,v,master_weights on host automatically in case of OOM. will run slower but won't OOM --- llmc/cuda_common.h | 2 +- llmc/cuda_utils.cuh | 19 ++++++++----------- train_gpt2.cu | 21 ++++++++++++++++----- 3 files changed, 25 insertions(+), 17 deletions(-) diff --git a/llmc/cuda_common.h b/llmc/cuda_common.h index 2b07aab..6f5bf65 100644 --- a/llmc/cuda_common.h +++ b/llmc/cuda_common.h @@ -48,7 +48,7 @@ constexpr std::bool_constant False; // ---------------------------------------------------------------------------- // Error checking -// CUDA error checking +// CUDA error checking. Underscore added so this function can be called directly not just via macro inline void cudaCheck_(cudaError_t error, const char *file, int line) { if (error != cudaSuccess) { printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, cudaGetErrorString(error)); diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index d3df74d..030ec07 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -208,23 +208,20 @@ void global_sum_deterministic(float* result, const Float* values, int count, cud // ---------------------------------------------------------------------------- // memory management -// allocate memory, preferrably on the -void cudaMallocConditionallyManaged(void** out, size_t bytes, const char *file, int line) { - // try to allocate `bytes` on device +// allocate memory, preferrably on the device +// returns a status code. 0 = OK, 1 = fell back to managed memory +int cudaMallocConditionallyManaged(void** out, size_t bytes, const char *file, int line) { + // try to allocate cudaError_t err = cudaMalloc(out, bytes); if(err == cudaErrorMemoryAllocation) { - // if that fails, fallback to a managed allocation. It will be slower, but at least - // it won't crash. - fprintf(stderr, "[WARN] Not enough space to allocate %zu MiB on device.\n" - " Falling back to managed allocation.\n" - " Speed may be negatively affected.\n", - bytes / 1024 / 1024); - // reset the error before the next API call - cudaGetLastError(); + // if we OOM, fallback to a managed allocation. slower but at least won't crash. + cudaGetLastError(); // reset the error before the next API call cudaCheck_(cudaMallocManaged(out, bytes), file, line); cudaCheck_(cudaMemAdvise(*out, bytes, cudaMemAdviseSetPreferredLocation, cudaCpuDeviceId), file, line); + return 1; } else { cudaCheck_(err, file, line); + return 0; } } diff --git a/train_gpt2.cu b/train_gpt2.cu index 343343e..289fe62 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -37,7 +37,7 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage. #include "llmc/cuda_common.h" // defines: // Packed128, f128, x128 -// warpReduceSum, warpReduceMax, blockReduce, copy_and_cast_kernel +// warpReduceSum, warpReduceMax, blockReduce, copy_and_cast_kernel, cudaMallocConditionallyManaged #include "llmc/cuda_utils.cuh" // defines: CUBLAS_LOWP, cublasCheck, cublaslt_workspace_size, cublaslt_workspace // defines: cublas_compute, cublaslt_handle, cublas_handle @@ -388,24 +388,35 @@ void gpt2_allocate_state(GPT2 *model, int B, int T) { model->workload_indices = (int*)mallocCheck(sizeof(int) * model->batch_size * model->seq_len * num_c_groups); model->bucket_info = (int4*)mallocCheck(sizeof(int4) * model->batch_size * model->seq_len * num_c_groups); + // cudaMallocConditionallyManaged can fall back to cudaMallocManaged if not enough memory on device + // and returns a status code of 1 if it had to fall back, in that case we want to print warning. + int memory_status = 0; + + // we will now init the optimizer states and master weights + // this is usually a substantial amount of memory allocation right here. size_t shard_num_parameters = multi_gpu_config.shard_num_parameters; // num parameters we are responsible for printf0("allocating %zu MiB for AdamW optimizer state m\n", (shard_num_parameters * sizeof(float)) >> 20); printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); assert(model->m_memory == nullptr); assert(model->v_memory == nullptr); - cudaMallocConditionallyManaged((void**)&model->m_memory, shard_num_parameters * sizeof(float)); - cudaMallocConditionallyManaged((void**)&model->v_memory, shard_num_parameters * sizeof(float)); + memory_status |= cudaMallocConditionallyManaged((void**)&model->m_memory, shard_num_parameters * sizeof(float)); + memory_status |= cudaMallocConditionallyManaged((void**)&model->v_memory, shard_num_parameters * sizeof(float)); if (model->use_master_weights == 1) { assert(model->master_weights == nullptr); printf0("allocating %zu MiB for master copy of params\n", (shard_num_parameters * sizeof(float)) >> 20); - cudaMallocConditionallyManaged((void**) &model->master_weights, shard_num_parameters * sizeof(float)); + memory_status |= cudaMallocConditionallyManaged((void**) &model->master_weights, shard_num_parameters * sizeof(float)); } + // report on mixed memory allocation status + if (memory_status == 1) { + printf0("WARNING: fell back to cudaMallocManaged when initializing m,v,master_weights.\n"); + printf0(" Prevents an OOM, but code may run much slower due to device <-> host memory movement\n"); + } + // report on device memory usage size_t free, total; cudaCheck(cudaMemGetInfo(&free, &total)); printf0("device memory usage: %zd MiB / %zd MiB\n", (total-free) / 1024 / 1024, total / 1024 / 1024); - // give an estimate of the maximum batch size size_t bytes_per_sequence = 0; for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) { From 8c586f915808a208373c76456f1cad84de66b4c9 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Aug 2024 18:08:34 +0000 Subject: [PATCH 6/7] reduce across GPUs nicer --- train_gpt2.cu | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 289fe62..809eca1 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -408,9 +408,10 @@ void gpt2_allocate_state(GPT2 *model, int B, int T) { memory_status |= cudaMallocConditionallyManaged((void**) &model->master_weights, shard_num_parameters * sizeof(float)); } - // report on mixed memory allocation status - if (memory_status == 1) { - printf0("WARNING: fell back to cudaMallocManaged when initializing m,v,master_weights.\n"); + // report on mixed memory allocation status (re-using our float reduce function, bit awk ok) + int readuced_memory_status = (int) multi_gpu_cpu_float_sum((float)memory_status, &multi_gpu_config); + if (readuced_memory_status >= 1) { + printf0("WARNING: Fell back to cudaMallocManaged when initializing m,v,master_weights on %d GPUs\n", readuced_memory_status); printf0(" Prevents an OOM, but code may run much slower due to device <-> host memory movement\n"); } // report on device memory usage From 18298f3a408b666557133628b30dca2135cc3775 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Aug 2024 18:16:46 +0000 Subject: [PATCH 7/7] i misspelled reduced --- train_gpt2.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 809eca1..70d8d0c 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -409,9 +409,9 @@ void gpt2_allocate_state(GPT2 *model, int B, int T) { } // report on mixed memory allocation status (re-using our float reduce function, bit awk ok) - int readuced_memory_status = (int) multi_gpu_cpu_float_sum((float)memory_status, &multi_gpu_config); - if (readuced_memory_status >= 1) { - printf0("WARNING: Fell back to cudaMallocManaged when initializing m,v,master_weights on %d GPUs\n", readuced_memory_status); + int reduced_memory_status = (int) multi_gpu_cpu_float_sum((float)memory_status, &multi_gpu_config); + if (reduced_memory_status >= 1) { + printf0("WARNING: Fell back to cudaMallocManaged when initializing m,v,master_weights on %d GPUs\n", reduced_memory_status); printf0(" Prevents an OOM, but code may run much slower due to device <-> host memory movement\n"); } // report on device memory usage