fix bug with all_hostsname_hashes missing, re-arrange code a bit, make sure we only print on rank0

This commit is contained in:
Andrej Karpathy 2024-04-26 18:11:09 +00:00
parent 2e17140d33
commit a00a3a88ef
4 changed files with 152 additions and 136 deletions

View file

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

View file

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

View file

@ -18,6 +18,7 @@ and fp8 (coming soon^TM).
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <ctype.h>
#include <math.h>
#include <time.h>
@ -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);

View file

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