aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-27 20:25:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fall back to cudaMallocManaged for optimizer states if we're out of memory

Browse source
This commit is contained in:
Erik Schultheis 2024-07-24 15:46:30 +02:00
parent 4c84bc743c
commit 0d52d2a3d7
3 changed files with 28 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

4
llmc/cuda_common.h
View file

@ -49,13 +49,13 @@ constexpr std::bool_constant<true> 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<class T>

23
llmc/cuda_utils.cuh
View file

@ -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

6
train_gpt2.cu
View file

@ -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;