[Multi-GPU] llm.c now runs on multiple GPUs with NCCL

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

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 <Number of GPUs on your machine> ./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:

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 <assert.h>
+#include <cuda_runtime.h>
+#include <mpi.h>
+#include <nccl.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+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<<<n_blocks, threads_per_block>>>(
+      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);
+}

train_gpt2.cu

@@ -32,6 +32,11 @@ and fp8 (coming soon^TM).
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 
+#ifdef MULTI_GPU
+#include <mpi.h>
+#include <nccl.h>
+#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<uintptr_t>(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;
 }

train_gpt2_fp32.cu

@@ -27,6 +27,11 @@ the layernorms are connected to the residuals so we += in layernorm backward.
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 
+#ifdef MULTI_GPU
+#include <mpi.h>
+#include <nccl.h>
+#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;
 }