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

move set_zero_configs into zero.cuh

Browse source
This commit is contained in:
Erik Schultheis 2024-07-09 17:12:37 +02:00
parent bdb0fb5599
commit 09b78bc902
2 changed files with 31 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

32
llmc/zero.cuh
View file

@ -18,7 +18,7 @@ Utilities for ZeRO sharding
#endif
#endif
// defines: fcloseCheck, fwriteCheck, scloseCheck, sclosesocketCheckCheck
// defines: fcloseCheck, fwriteCheck, scloseCheck, sclosesocketCheck
#include "utils.h"
// ----------------------------------------------------------------------------
@ -78,6 +78,10 @@ typedef struct {
#endif
} MultiGpuConfig;
// one global variable to hold the multi-GPU configuration for this process
// inline, so we can include this header multiple times without getting multiple definitions
inline MultiGpuConfig multi_gpu_config;
#ifdef MULTI_GPU
#ifdef _WIN32
@ -548,5 +552,31 @@ void multi_gpu_async_reduce_gradient(
#endif
}
// convenience macro that only prints if the rank of process is zero
#define printf0(...) if (::multi_gpu_config.process_rank == 0) { printf(__VA_ARGS__); }
void set_zero_configs(MultiGpuConfig* multi_gpu_config, int zero_stage, size_t total_parameters) {
multi_gpu_config->zero_stage = 0;
multi_gpu_config->shard_num_parameters = total_parameters;
// Check the Zero Stage and define sharding parameters
if (zero_stage == 0) {
printf0("| Zero Optimization is disabled |\n");
}
else if (zero_stage == 1) {
if (total_parameters % multi_gpu_config->num_processes != 0) {
printf0("| Zero Optimization is disabled, Can't equally partition parameters |\n");
multi_gpu_config->zero_stage = 0;
}
else {
multi_gpu_config->zero_stage = 1;
multi_gpu_config->shard_num_parameters = total_parameters / multi_gpu_config->num_processes;
}
}
else{
printf0("| Disabling Zero Optimization, Zero Stage2 and Stage3 are not yet supported |\n");
multi_gpu_config->zero_stage = 0;
}
}
#endif

35
train_gpt2.cu
View file

@ -73,44 +73,9 @@ char filename_buffer[512];
// global vars containing information about the GPU this process is running on
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, ...) {
if (multi_gpu_config.process_rank == 0) {
va_list args;
va_start(args, format);
vprintf(format, args);
va_end(args);
}
}
void set_zero_configs(MultiGpuConfig* multi_gpu_config, int zero_stage, size_t total_parameters) {
multi_gpu_config->zero_stage = 0;
multi_gpu_config->shard_num_parameters = total_parameters;
// Check the Zero Stage and define sharding parameters
if (zero_stage == 0) {
printf0("| Zero Optimization is disabled |\n");
}
else if (zero_stage == 1) {
if (total_parameters % multi_gpu_config->num_processes != 0) {
printf0("| Zero Optimization is disabled, Can't equally partition parameters |\n");
multi_gpu_config->zero_stage = 0;
}
else {
multi_gpu_config->zero_stage = 1;
multi_gpu_config->shard_num_parameters = total_parameters / multi_gpu_config->num_processes;
}
}
else{
printf0("| Disabling Zero Optimization, Zero Stage2 and Stage3 are not yet supported |\n");
multi_gpu_config->zero_stage = 0;
}
}
// ----------------------------------------------------------------------------
// GPT-2 model definition