diff --git a/.ipynb_checkpoints/Makefile-checkpoint b/.ipynb_checkpoints/Makefile-checkpoint new file mode 100644 index 0000000..a50c642 --- /dev/null +++ b/.ipynb_checkpoints/Makefile-checkpoint @@ -0,0 +1,324 @@ +CC ?= clang +CFLAGS = -Ofast -Wno-unused-result -Wno-ignored-pragmas -Wno-unknown-attributes +LDFLAGS = +LDLIBS = -lm +INCLUDES = +CFLAGS_COND = -march=native + +# Find nvcc +SHELL_UNAME = $(shell uname) +REMOVE_FILES = rm -f +OUTPUT_FILE = -o $@ +CUDA_OUTPUT_FILE = -o $@ + +# Default O3 CPU optimization level for NVCC (0 for fastest compile time) +FORCE_NVCC_O ?= 3 + +# NVCC flags +# -t=0 is short for --threads, 0 = number of CPUs on the machine +NVCC_FLAGS = --threads=0 -t=0 --use_fast_math -std=c++17 -O$(FORCE_NVCC_O) +NVCC_LDFLAGS = -lcublas -lcublasLt +NVCC_INCLUDES = +NVCC_LDLIBS = +NCLL_INCUDES = +NVCC_CUDNN = +# By default we don't build with cudnn because it blows up compile time from a few seconds to ~minute +USE_CUDNN ?= 0 + +# We will place .o files in the `build` directory (create it if it doesn't exist) +BUILD_DIR = build +ifeq ($(OS), Windows_NT) + $(shell if not exist $(BUILD_DIR) mkdir $(BUILD_DIR)) + REMOVE_BUILD_OBJECT_FILES := del $(BUILD_DIR)\*.obj +else + $(shell mkdir -p $(BUILD_DIR)) + REMOVE_BUILD_OBJECT_FILES := rm -f $(BUILD_DIR)/*.o +endif + +# Function to check if a file exists in the PATH +ifneq ($(OS), Windows_NT) +define file_exists_in_path + $(which $(1) 2>/dev/null) +endef +else +define file_exists_in_path + $(shell where $(1) 2>nul) +endef +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),) + # 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 +endif + +# set to defaults if - make GPU_COMPUTE_CAPABILITY= otherwise use the compute capability detected above +ifneq ($(GPU_COMPUTE_CAPABILITY),) + NVCC_FLAGS += --generate-code arch=compute_$(GPU_COMPUTE_CAPABILITY),code=[compute_$(GPU_COMPUTE_CAPABILITY),sm_$(GPU_COMPUTE_CAPABILITY)] +endif + +# autodect a lot of various supports on current platform +$(info ---------------------------------------------) + +ifneq ($(OS), Windows_NT) + NVCC := $(shell which nvcc 2>/dev/null) + NVCC_LDFLAGS += -lnvidia-ml + + # Function to test if the compiler accepts a given flag. + define check_and_add_flag + $(eval FLAG_SUPPORTED := $(shell printf "int main() { return 0; }\n" | $(CC) $(1) -x c - -o /dev/null 2>/dev/null && echo 'yes')) + ifeq ($(FLAG_SUPPORTED),yes) + CFLAGS += $(1) + endif + endef + + # Check each flag and add it if supported + $(foreach flag,$(CFLAGS_COND),$(eval $(call check_and_add_flag,$(flag)))) +else + CFLAGS := + REMOVE_FILES = del *.exe,*.obj,*.lib,*.exp,*.pdb && del + SHELL_UNAME := Windows + ifneq ($(shell where nvcc 2> nul),"") + NVCC := nvcc + else + NVCC := + endif + CC := cl + CFLAGS = /Idev /Zi /nologo /W4 /WX- /diagnostics:column /sdl /O2 /Oi /Ot /GL /D _DEBUG /D _CONSOLE /D _UNICODE /D UNICODE /Gm- /EHsc /MD /GS /Gy /fp:fast /Zc:wchar_t /Zc:forScope /Zc:inline /permissive- \ + /external:W3 /Gd /TP /wd4996 /Fd$@.pdb /FC /openmp:llvm + LDFLAGS := + LDLIBS := + INCLUDES := + NVCC_FLAGS += -I"dev" + ifeq ($(WIN_CI_BUILD),1) + $(info Windows CI build) + OUTPUT_FILE = /link /OUT:$@ + CUDA_OUTPUT_FILE = -o $@ + else + $(info Windows local build) + OUTPUT_FILE = /link /OUT:$@ && copy /Y $@ $@.exe + CUDA_OUTPUT_FILE = -o $@ && copy /Y $@.exe $@ + endif +endif + +# Check and include cudnn if available +# You can override the path to cudnn frontend by setting CUDNN_FRONTEND_PATH on the make command line +# By default, we look for it in HOME/cudnn-frontend/include and ./cudnn-frontend/include +# Refer to the README for cuDNN install instructions +ifeq ($(USE_CUDNN), 1) + ifeq ($(SHELL_UNAME), Linux) + ifeq ($(shell [ -d $(HOME)/cudnn-frontend/include ] && echo "exists"), exists) + $(info ✓ cuDNN found, will run with flash-attention) + CUDNN_FRONTEND_PATH ?= $(HOME)/cudnn-frontend/include + else ifeq ($(shell [ -d cudnn-frontend/include ] && echo "exists"), exists) + $(info ✓ cuDNN found, will run with flash-attention) + CUDNN_FRONTEND_PATH ?= cudnn-frontend/include + else + $(error ✗ cuDNN not found. See the README for install instructions and the Makefile for hard-coded paths) + endif + NVCC_INCLUDES += -I$(CUDNN_FRONTEND_PATH) + NVCC_LDFLAGS += -lcudnn + NVCC_FLAGS += -DENABLE_CUDNN + NVCC_CUDNN = $(BUILD_DIR)/cudnn_att.o + else + ifneq ($(OS), Windows_NT) + $(info → cuDNN is not supported on MAC OS right now) + else + $(info ✓ Windows cuDNN found, will run with flash-attention) + ifeq ($(shell if exist "$(HOMEDRIVE)$(HOMEPATH)\cudnn-frontend\include" (echo exists)),exists) + CUDNN_FRONTEND_PATH ?= $(HOMEDRIVE)$(HOMEPATH)\cudnn-frontend\include #override on command line if different location + else ifeq ($(shell if exist "cudnn-frontend\include" (echo exists)),exists) + CUDNN_FRONTEND_PATH ?= cudnn-frontend\include #override on command line if different location + else + $(error ✗ cuDNN not found. See the README for install instructions and the Makefile for hard-coded paths) + endif + CUDNN_INCLUDE_PATH ?= -I"C:\Program Files\NVIDIA\CUDNN\v9.1\include\12.4" + CUDNN_FRONTEND_PATH += $(CUDNN_INCLUDE_PATH) + NVCC_FLAGS += --std c++20 -Xcompiler "/std:c++20" -Xcompiler "/EHsc /W0 /nologo /Ox /FS" -maxrregcount=0 --machine 64 + NVCC_CUDNN = $(BUILD_DIR)\cudnn_att.obj + NVCC_INCLUDES += -I$(CUDNN_FRONTEND_PATH) + NVCC_LDFLAGS += -L"C:\Program Files\NVIDIA\CUDNN\v9.1\lib\12.4\x64" -lcudnn + NVCC_FLAGS += -DENABLE_CUDNN + endif + endif +else + $(info → cuDNN is manually disabled by default, run make with `USE_CUDNN=1` to try to enable) +endif + +# Check if OpenMP is available +# This is done by attempting to compile an empty file with OpenMP flags +# OpenMP makes the code a lot faster so I advise installing it +# e.g. on MacOS: brew install libomp +# e.g. on Ubuntu: sudo apt-get install libomp-dev +# later, run the program by prepending the number of threads, e.g.: OMP_NUM_THREADS=8 ./gpt2 +# First, check if NO_OMP is set to 1, if not, proceed with the OpenMP checks +ifeq ($(NO_OMP), 1) + $(info OpenMP is manually disabled) +else + ifneq ($(OS), Windows_NT) + # Detect if running on macOS or Linux + ifeq ($(SHELL_UNAME), Darwin) + # Check for Homebrew's libomp installation in different common directories + ifeq ($(shell [ -d /opt/homebrew/opt/libomp/lib ] && echo "exists"), exists) + # macOS with Homebrew on ARM (Apple Silicon) + CFLAGS += -Xclang -fopenmp -DOMP + LDFLAGS += -L/opt/homebrew/opt/libomp/lib + LDLIBS += -lomp + INCLUDES += -I/opt/homebrew/opt/libomp/include + $(info ✓ OpenMP found) + else ifeq ($(shell [ -d /usr/local/opt/libomp/lib ] && echo "exists"), exists) + # macOS with Homebrew on Intel + CFLAGS += -Xclang -fopenmp -DOMP + LDFLAGS += -L/usr/local/opt/libomp/lib + LDLIBS += -lomp + INCLUDES += -I/usr/local/opt/libomp/include + $(info ✓ OpenMP found) + else + $(info ✗ OpenMP not found) + endif + else + # Check for OpenMP support in GCC or Clang on Linux + ifeq ($(shell echo | $(CC) -fopenmp -x c -E - > /dev/null 2>&1; echo $$?), 0) + CFLAGS += -fopenmp -DOMP + LDLIBS += -lgomp + $(info ✓ OpenMP found) + else + $(info ✗ OpenMP not found) + endif + endif + endif +endif + +# Check if NCCL is available, include if so, for multi-GPU training +ifeq ($(NO_MULTI_GPU), 1) + $(info → Multi-GPU (NCCL) is manually disabled) +else + ifneq ($(OS), Windows_NT) + # Detect if running on macOS or Linux + ifeq ($(SHELL_UNAME), Darwin) + $(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping NCCL support) + else ifeq ($(shell dpkg -l | grep -q nccl && echo "exists"), exists) + $(info ✓ NCCL found, OK to train with multiple GPUs) + NVCC_FLAGS += -DMULTI_GPU + NVCC_LDLIBS += -lnccl + else + $(info ✗ NCCL is not found, disabling multi-GPU support) + $(info ---> On Linux you can try install NCCL with `sudo apt install libnccl2 libnccl-dev`) + endif + endif +endif + +# Attempt to find and include OpenMPI on the system +OPENMPI_DIR ?= /usr/lib/x86_64-linux-gnu/openmpi +OPENMPI_LIB_PATH = $(OPENMPI_DIR)/lib/ +OPENMPI_INCLUDE_PATH = $(OPENMPI_DIR)/include/ +ifeq ($(NO_USE_MPI), 1) + $(info → MPI is manually disabled) +else ifeq ($(shell [ -d $(OPENMPI_LIB_PATH) ] && [ -d $(OPENMPI_INCLUDE_PATH) ] && echo "exists"), exists) + $(info ✓ MPI enabled) + NVCC_INCLUDES += -I$(OPENMPI_INCLUDE_PATH) + NVCC_LDFLAGS += -L$(OPENMPI_LIB_PATH) + NVCC_LDLIBS += -lmpi + NVCC_FLAGS += -DUSE_MPI +else + $(info ✗ MPI not found) +endif + +# Precision settings, default to bf16 but ability to override +PRECISION ?= BF16 +VALID_PRECISIONS := FP32 FP16 BF16 +ifeq ($(filter $(PRECISION),$(VALID_PRECISIONS)),) + $(error Invalid precision $(PRECISION), valid precisions are $(VALID_PRECISIONS)) +endif +ifeq ($(PRECISION), FP32) + PFLAGS = -DENABLE_FP32 +else ifeq ($(PRECISION), FP16) + PFLAGS = -DENABLE_FP16 +else + PFLAGS = -DENABLE_BF16 +endif + +# PHONY means these targets will always be executed +.PHONY: all train_gpt2 test_gpt2 train_gpt2cu test_gpt2cu train_gpt2fp32cu test_gpt2fp32cu profile_gpt2cu + +# Add targets +TARGETS = train_gpt2 test_gpt2 + +# Conditional inclusion of CUDA targets +ifeq ($(NVCC),) + $(info ✗ nvcc not found, skipping GPU/CUDA builds) +else + $(info ✓ nvcc found, including GPU/CUDA support) + TARGETS += train_gpt2cu test_gpt2cu train_gpt2fp32cu test_gpt2fp32cu $(NVCC_CUDNN) +endif + +$(info ---------------------------------------------) + +all: $(TARGETS) + +train_gpt2: train_gpt2.c + $(CC) $(CFLAGS) $(INCLUDES) $(LDFLAGS) $^ $(LDLIBS) $(OUTPUT_FILE) + +test_gpt2: test_gpt2.c + $(CC) $(CFLAGS) $(INCLUDES) $(LDFLAGS) $^ $(LDLIBS) $(OUTPUT_FILE) + +$(NVCC_CUDNN): llmc/cudnn_att.cpp + $(NVCC) -c $(NVCC_FLAGS) $(PFLAGS) $^ $(NVCC_INCLUDES) -o $@ + +train_gpt2cu: train_gpt2.cu $(NVCC_CUDNN) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) + +train_gpt2fp32cu: train_gpt2_fp32.cu + $(NVCC) $(NVCC_FLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) + +test_gpt2cu: test_gpt2.cu $(NVCC_CUDNN) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) + +test_gpt2fp32cu: test_gpt2_fp32.cu + $(NVCC) $(NVCC_FLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) + +profile_gpt2cu: profile_gpt2.cu $(NVCC_CUDNN) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) + +# NCCL test target +test_nccl: test_nccl_basic.cu + $(NVCC) $(NVCC_FLAGS) $^ -lnccl -lcudart $(CUDA_OUTPUT_FILE) + +# NVSHMEM test target +# Run with: NVSHMEM_BOOTSTRAP=MPI mpirun -np 2 ./test_nvshmem +# NVSHMEM requires sm_80+ for system-scope atomics and extended C++ features +NVSHMEM_HOME ?= /usr/local/nvshmem +test_nvshmem: test_nvshmem_basic.cu + $(NVCC) --threads=0 -t=0 --use_fast_math -std=c++17 -O3 \ + -arch=sm_80 --extended-lambda --expt-relaxed-constexpr \ + -rdc=true -I$(NVSHMEM_HOME)/include -L$(NVSHMEM_HOME)/lib \ + -I$(OPENMPI_INCLUDE_PATH) -L$(OPENMPI_LIB_PATH) \ + $^ -lnvshmem_host -lnvshmem_device -lcuda -lnvidia-ml -lcudart -lmpi $(CUDA_OUTPUT_FILE) + +# NVSHMEM Pipeline Training Target (Phase 1) +# Run with: NVSHMEM_BOOTSTRAP=MPI mpirun -np 2 ./nvshmem_train_gpt2 -b 4 -t 128 +nvshmem_train_gpt2: nvshmem_train_gpt2.cu + $(NVCC) --threads=0 -t=0 --use_fast_math -std=c++14 -O3 \ + -arch=sm_80 \ + -I$(NVSHMEM_HOME)/include -L$(NVSHMEM_HOME)/lib \ + -I$(OPENMPI_INCLUDE_PATH) -L$(OPENMPI_LIB_PATH) \ + $^ -lnvshmem_host -lnccl -lnvshmem_device -lcublas -lcublasLt -lcuda -lcudart -lnvidia-ml -lmpi \ + $(if $(filter 1,$(NCCL_FOUND)),-lnccl,) \ + $(CUDA_OUTPUT_FILE) +nvshmem_train_gpt2_partitioned: nvshmem_train_gpt2_partitioned.cu + $(NVCC) --threads=0 -t=0 --use_fast_math -std=c++14 -O3 \ + -arch=sm_80 \ + -I$(NVSHMEM_HOME)/include -L$(NVSHMEM_HOME)/lib \ + -I$(OPENMPI_INCLUDE_PATH) -L$(OPENMPI_LIB_PATH) \ + $^ -lnvshmem_host -lnccl -lnvshmem_device -lcublas -lcublasLt -lcuda -lcudart -lnvidia-ml -lmpi \ + $(if $(filter 1,$(NCCL_FOUND)),-lnccl,) \ + $(CUDA_OUTPUT_FILE) + +clean: + $(REMOVE_FILES) $(TARGETS) test_nccl test_nvshmem nvshmem_train_gpt2 + $(REMOVE_BUILD_OBJECT_FILES) diff --git a/.ipynb_checkpoints/nvshmem_train_gpt2-checkpoint.cu b/.ipynb_checkpoints/nvshmem_train_gpt2-checkpoint.cu index b49432c..816c8ca 100644 --- a/.ipynb_checkpoints/nvshmem_train_gpt2-checkpoint.cu +++ b/.ipynb_checkpoints/nvshmem_train_gpt2-checkpoint.cu @@ -1216,7 +1216,7 @@ typedef struct { } GPT2; void gpt2_build_from_checkpoint(GPT2 *model, const char *checkpoint_path) { - + printf("oioioioioi_anand\n"); // read in model from a checkpoint file FILE *model_file = fopenCheck(checkpoint_path, "rb"); int model_header[256]; @@ -1338,7 +1338,6 @@ void printer(int B, int T, int C, float* arr, char* initial) void gpt2_forward(GPT2 *model, int *inputs, int *targets, int B, int T) { // targets are optional and could be NULL - counter++; // ensure the model was initialized or error out if (model->params_memory == NULL) { printf("Error: model was not initialized properly.\n"); @@ -1468,6 +1467,7 @@ void gpt2_forward(GPT2 *model, int *inputs, int *targets, int B, int T) { // Use nvshmem_putmem to transfer directly to PE 1 + cudaCheck(cudaDeviceSynchronize()); nvshmem_putmem(model->nvshmem_act_buffer, layer5_output, B * T * C * sizeof(float), 1); // PE 1 nvshmem_quiet(); @@ -1589,6 +1589,8 @@ void gpt2_zero_grad(GPT2 *model) { void gpt2_backward(GPT2 *model) { + counter++; + // double check we forwarded previously, with targets if (model->mean_loss == -1.0f) { printf("Error: must forward with targets before backward\n"); @@ -1717,13 +1719,12 @@ void gpt2_backward(GPT2 *model) { l_ln1_mean, l_ln1_rstd, B, T, C); } + cudaCheck(cudaDeviceSynchronize()); // Transfer gradients to GPU 0 using nvshmem_putmem nvshmem_putmem(model->nvshmem_grad_buffer, dresidual, B * T * C * sizeof(float), 0); // PE 0 nvshmem_quiet(); - printer(B,T,C, model->nvshmem_grad_buffer, "BACKWARD:nvshmem_grad_buffer"); - printer(B,T,C, dresidual, "BACKWARD:dresidual"); } @@ -1732,16 +1733,10 @@ void gpt2_backward(GPT2 *model) { if (my_pe == 0) { // Wait for GPU 1 to send gradients + cudaMemcpy(dresidual, model->nvshmem_grad_buffer, B * T * C * sizeof(float), cudaMemcpyDeviceToDevice); - printer(B,T,C, model->nvshmem_grad_buffer, "BACKWARD:nvshmem_grad_buffer"); - printer(B,T,C, dresidual, "BACKWARD:dresidual"); - - // Copy received gradients from GPU 1's symmetric buffer - // nvshmem_getmem(dresidual, model->nvshmem_grad_buffer, - // B * T * C * sizeof(float), 1); // From PE 1 - // Backward through layers 5-0 for (int l = 5; l >= 0; l--) { residual = (l == 0) ? acts.encoded : acts.residual3 + (l - 1) * B * T * C; @@ -2276,9 +2271,11 @@ int main(int argc, char *argv[]) { (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; + if (my_pe == 1) { 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); + } logger_log_train(&logger, step, model.mean_loss); } // add a total average, for optimizations that are only mild improvements diff --git a/Makefile b/Makefile index 7c696f0..a50c642 100644 --- a/Makefile +++ b/Makefile @@ -310,6 +310,14 @@ nvshmem_train_gpt2: nvshmem_train_gpt2.cu $^ -lnvshmem_host -lnccl -lnvshmem_device -lcublas -lcublasLt -lcuda -lcudart -lnvidia-ml -lmpi \ $(if $(filter 1,$(NCCL_FOUND)),-lnccl,) \ $(CUDA_OUTPUT_FILE) +nvshmem_train_gpt2_partitioned: nvshmem_train_gpt2_partitioned.cu + $(NVCC) --threads=0 -t=0 --use_fast_math -std=c++14 -O3 \ + -arch=sm_80 \ + -I$(NVSHMEM_HOME)/include -L$(NVSHMEM_HOME)/lib \ + -I$(OPENMPI_INCLUDE_PATH) -L$(OPENMPI_LIB_PATH) \ + $^ -lnvshmem_host -lnccl -lnvshmem_device -lcublas -lcublasLt -lcuda -lcudart -lnvidia-ml -lmpi \ + $(if $(filter 1,$(NCCL_FOUND)),-lnccl,) \ + $(CUDA_OUTPUT_FILE) clean: $(REMOVE_FILES) $(TARGETS) test_nccl test_nvshmem nvshmem_train_gpt2 diff --git a/nvshmem_train_gpt2 b/nvshmem_train_gpt2 deleted file mode 100755 index 5989589..0000000 Binary files a/nvshmem_train_gpt2 and /dev/null differ diff --git a/nvshmem_train_gpt2.cu b/nvshmem_train_gpt2.cu index fd04b97..eed2cbc 100644 --- a/nvshmem_train_gpt2.cu +++ b/nvshmem_train_gpt2.cu @@ -1216,7 +1216,6 @@ typedef struct { } GPT2; void gpt2_build_from_checkpoint(GPT2 *model, const char *checkpoint_path) { - // read in model from a checkpoint file FILE *model_file = fopenCheck(checkpoint_path, "rb"); int model_header[256]; diff --git a/nvshmem_train_gpt2_partitioned.cu b/nvshmem_train_gpt2_partitioned.cu new file mode 100644 index 0000000..31fc486 --- /dev/null +++ b/nvshmem_train_gpt2_partitioned.cu @@ -0,0 +1,1977 @@ +/* +GPT-2 Transformer Neural Net trained in raw CUDA with Model Partitioning +This version partitions the model across GPUs to save memory. +Each GPU only holds its assigned layers' parameters and activations. +*/ + +#include +#include +#include +#include +#include +#include +#include +#include + +// GPU / CUDA related +#include +#include +#include +#include +// NVSHMEM for multi-GPU communication +#include +#include +#include +#include +// our own utilities +#include "llmc/utils.h" +#include "llmc/tokenizer.h" +#include "llmc/dataloader.h" + +// ---------------------------------------------------------------------------- +// CUDA utils + +#define CEIL_DIV(M, N) (((M) + (N) - 1) / (N)) + +int counter = 0; + +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__)) + +void cublasCheck(cublasStatus_t status, const char *file, int line) { + if (status != CUBLAS_STATUS_SUCCESS) { + printf("[cuBLAS ERROR]: %d %s %d\n", status, file, line); + exit(EXIT_FAILURE); + } +} +#define cublasCheck(status) { cublasCheck((status), __FILE__, __LINE__); } + +void ncclCheck(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) (ncclCheck(err, __FILE__, __LINE__)) + +static cublasComputeType_t cublas_compute_type; +cublasHandle_t cublas_handle; + +namespace cg = cooperative_groups; + +// ---------------------------------------------------------------------------- +// Pipeline Partition Structure + +typedef struct { + int first_layer; // First layer index (global) + int num_layers; // Number of layers this PE handles + int has_embedding; // 1 if this PE does embedding (PE 0) + int has_final_ln; // 1 if this PE does final layernorm (PE 1) +} PipelinePartition; + +void get_partition_for_pe(PipelinePartition* part, int my_pe, int n_pes, int total_layers) { + int layers_per_pe = total_layers / n_pes; + part->first_layer = my_pe * layers_per_pe; + part->num_layers = layers_per_pe; + part->has_embedding = (my_pe == 0); + part->has_final_ln = (my_pe == n_pes - 1); +} + +// ---------------------------------------------------------------------------- +// all the kernels + +__device__ inline float4 add_float4(const float4 &a, const float4 &b) { + return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); +} + +__global__ void encoder_forward_kernel3(float4 *out, const int *inp, + const float4 *wte, const float4 *wpe, + int B, int T, int C) { + int C4 = C / 4; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int N = B * T * C4; + if (idx < N) { + int bt = idx / C4; + int b = bt / T; + int t = bt % T; + int c4 = idx % C4; + int ix = inp[b * T + t]; + out[b * T * C4 + t * C4 + c4] = + add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]); + } +} + +__global__ void encoder_backward_kernel(float *dwte, float *dwpe, + const float *dout, const int *inp, + int B, int T, int C) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int N = B * T * C; + if (idx < N) { + int bt = idx / C; + int b = bt / T; + int t = bt % T; + int c = idx % C; + int ix = inp[b * T + t]; + const float *dout_btc = dout + b * T * C + t * C + c; + float *dwte_ix = dwte + ix * C + c; + float *dwpe_tc = dwpe + t * C + c; + atomicAdd(dwte_ix, *dout_btc); + atomicAdd(dwpe_tc, *dout_btc); + } +} + +__global__ void layernorm_forward_kernel3( + float *__restrict__ out, float *__restrict__ mean, float *__restrict__ rstd, + const float *__restrict__ inp, const float *__restrict__ weight, + const float *__restrict__ bias, int N, int C) { + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); + if (idx >= N) return; + + const float *x = inp + idx * C; + float sum = 0.0f; + for (int i = warp.thread_rank(); i < C; i += warp.size()) { + sum += x[i]; + } + sum = cg::reduce(warp, sum, cg::plus{}); + float m = sum / C; + if (warp.thread_rank() == 0 && mean != nullptr) { + __stcs(mean + idx, m); + } + + sum = 0.0f; + for (int i = warp.thread_rank(); i < C; i += warp.size()) { + float diff = x[i] - m; + sum += diff * diff; + } + sum = cg::reduce(warp, sum, cg::plus{}); + float s = rsqrtf(sum / C + 1e-5f); + if (warp.thread_rank() == 0 && rstd != nullptr) { + __stcs(rstd + idx, s); + } + + float *o = out + idx * C; + for (int c = warp.thread_rank(); c < C; c += warp.size()) { + float n = s * (__ldcs(x + c) - m); + __stcs(o + c, n * weight[c] + bias[c]); + } +} + +__global__ void permute_kernel(float *q, float *k, float *v, const float *inp, + int B, int N, int NH, int d) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < B * NH * N * d) { + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; + q[idx] = __ldcs(&inp[inp_idx]); + k[idx] = __ldcs(&inp[inp_idx + NH * d]); + v[idx] = __ldcs(&inp[inp_idx + 2 * (NH * d)]); + } +} + +__global__ void permute_kernel_backward(float *dinp, const float *dq, + const float *dk, const float *dv, int B, + int N, int NH, int d) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < B * NH * N * d) { + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; + dinp[inp_idx] = dq[idx]; + dinp[inp_idx + NH * d] = dk[idx]; + dinp[inp_idx + 2 * (NH * d)] = dv[idx]; + } +} + +__global__ void unpermute_kernel(float *inp, float *out, int B, int N, int NH, int d) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < B * NH * N * d) { + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; + out[other_idx] = __ldcs(&inp[idx]); + } +} + +__global__ void unpermute_kernel_backward(float *dinp, const float *dout, int B, + int N, int NH, int d) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < B * NH * N * d) { + int b = idx / (NH * N * d); + int rest = idx % (NH * N * d); + int nh_ = rest / (N * d); + rest = rest % (N * d); + int n = rest / d; + int d_ = rest % d; + int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; + dinp[idx] = dout[other_idx]; + } +} + +__device__ float &vec_at(float4 &vec, int index) { + return reinterpret_cast(&vec)[index]; +} + +__device__ float vec_at(const float4 &vec, int index) { + return reinterpret_cast(&vec)[index]; +} + +__global__ void softmax_forward_kernel5(float *out, float inv_temperature, + const float *inp, int N, int T) { + assert(T % 4 == 0); + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + int idx = (gridDim.x - blockIdx.x - 1) * warp.meta_group_size() + warp.meta_group_rank(); + if (idx >= N * T) return; + int own_pos = idx % T; + int pos_by_4 = own_pos / 4; + + const float *x = inp + idx * T; + float maxval = -FLT_MAX; + float sumval = 0.0f; + + const float4 *x_vec = reinterpret_cast(x); + for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) { + float4 v = x_vec[i]; + float old_maxval = maxval; + for (int k = 0; k < 4; ++k) { + maxval = fmaxf(maxval, vec_at(v, k)); + } + sumval *= expf(inv_temperature * (old_maxval - maxval)); + for (int k = 0; k < 4; ++k) { + sumval += expf(inv_temperature * (vec_at(v, k) - maxval)); + } + } + + if (4 * pos_by_4 + warp.thread_rank() <= own_pos) { + float old_maxval = maxval; + maxval = fmaxf(maxval, x[4 * pos_by_4 + warp.thread_rank()]); + sumval *= expf(inv_temperature * (old_maxval - maxval)); + sumval += expf(inv_temperature * (x[4 * pos_by_4 + warp.thread_rank()] - maxval)); + } + + float global_maxval = cg::reduce(warp, maxval, cg::greater{}); + sumval *= expf(inv_temperature * (maxval - global_maxval)); + float sum = cg::reduce(warp, sumval, cg::plus{}); + float norm = 1.f / sum; + + for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { + float ev = expf(inv_temperature * (__ldcs(x + i) - global_maxval)); + __stcs(out + idx * T + i, ev * norm); + } +} + +__global__ void residual_forward_kernel(float *out, float *inp1, float *inp2, int N) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + out[idx] = __ldcs(&inp1[idx]) + __ldcs(&inp2[idx]); + } +} + +#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI) +__global__ void gelu_forward_kernel(float *out, const float *inp, int N) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < N) { + float xi = inp[i]; + float cube = 0.044715f * xi * xi * xi; + out[i] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))); + } +} + +__global__ void gelu_backward_kernel(float *dinp, const float *inp, + const float *dout, const int N) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < N) { + float x = inp[i]; + float cube = 0.044715f * x * x * x; + float tanh_arg = GELU_SCALING_FACTOR * (x + cube); + float tanh_out = tanhf(tanh_arg); + float coshf_out = coshf(tanh_arg); + float sech_out = 1.0f / (coshf_out * coshf_out); + float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * + (1.0f + 3.0f * 0.044715f * x * x); + dinp[i] = local_grad * dout[i]; + } +} + +__global__ void matmul_backward_bias_kernel4(float *dbias, const float *dout, + int B, int T, int OC) { + extern __shared__ float smem[]; + const int warp_id = threadIdx.x / warpSize; + const int lane_id = threadIdx.x % warpSize; + const int tl = blockIdx.x * warpSize; + const int vstep = blockDim.x / warpSize; + + const float *dout_col = dout + tl + lane_id; + float dout_sum = 0.0f; + for (int row = warp_id; row < B * T; row += vstep) { + dout_sum += dout_col[row * OC]; + } + smem[lane_id + warp_id * warpSize] = dout_sum; + __syncthreads(); + + dout_sum = 0.0f; + if (warp_id == 0) { + for (int j = 0; j < vstep; j++) { + dout_sum += smem[lane_id + j * warpSize]; + } + dbias[tl + lane_id] += dout_sum; + } +} + +__global__ void layernorm_backward_kernel2(float *dinp, float *dweight, + float *dbias, const float *dout, + const float *inp, const float *weight, + const float *mean, const float *rstd, + int B, int T, int C) { + extern __shared__ float shared[]; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); + int N = B * T; + if (idx >= N) return; + + int b = idx / T; + int t = idx % T; + + const float *dout_bt = dout + b * T * C + t * C; + const float *inp_bt = inp + b * T * C + t * C; + float *dinp_bt = dinp + b * T * C + t * C; + const float mean_bt = mean[b * T + t]; + const float rstd_bt = rstd[b * T + t]; + + float *dbias_shared = shared; + float *dweight_shared = shared + C; + + for (int i = threadIdx.x; i < C; i += blockDim.x) { + dbias_shared[i] = 0.0f; + dweight_shared[i] = 0.0f; + } + __syncthreads(); + + float dnorm_mean = 0.0f; + float dnorm_norm_mean = 0.0f; + for (int i = warp.thread_rank(); i < C; i += warp.size()) { + float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; + float dnorm_i = weight[i] * dout_bt[i]; + dnorm_mean += dnorm_i; + dnorm_norm_mean += dnorm_i * norm_bti; + } + dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus{}); + dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus{}); + dnorm_mean = dnorm_mean / C; + dnorm_norm_mean = dnorm_norm_mean / C; + + for (int i = warp.thread_rank(); i < C; i += warp.size()) { + float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; + float dnorm_i = weight[i] * dout_bt[i]; + atomicAdd(&dbias_shared[i], dout_bt[i]); + atomicAdd(&dweight_shared[i], norm_bti * dout_bt[i]); + float dval = 0.0f; + dval += dnorm_i; + dval -= dnorm_mean; + dval -= norm_bti * dnorm_norm_mean; + dval *= rstd_bt; + dinp_bt[i] += dval; + } + __syncthreads(); + + for (int i = threadIdx.x; i < C; i += blockDim.x) { + atomicAdd(&dbias[i], dbias_shared[i]); + atomicAdd(&dweight[i], dweight_shared[i]); + } +} + +__global__ void softmax_autoregressive_backward_kernel(float *dpreatt, + const float *datt, + const float *att, int B, + int T, int C, float scale) { + constexpr const int BlockSize = 256; + constexpr int T_per_block = 4; + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + __shared__ float block_acc[32]; + + int idx = blockIdx.y; + int t0 = T - 1 - T_per_block * blockIdx.x; + + att += idx * T * T; + datt += idx * T * T; + dpreatt += idx * T * T; + + if (warp.meta_group_rank() == 0) { + block_acc[warp.thread_rank()] = 0; + } + + for (int to = 0; to < T_per_block; ++to) { + int t = t0 - to; + if (t < 0) return; + const float *att_bth = att + t * T; + const float *datt_bth = datt + t * T; + float *dpreatt_bth = dpreatt + t * T; + + float local_sum = 0; + for (int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { + local_sum += att_bth[t2] * datt_bth[t2]; + } + + block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{}); + block.sync(); + local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{}); + + for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) { + float acc = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum); + __stcs(dpreatt_bth + t3, scale * acc); + } + } +} + +__device__ inline float lerp(float start, float end, float weight) { + return fma(weight, end, fma(-weight, start, start)); +} + +__global__ void adamw_kernel2(float *params_memory, float *grads_memory, + float *m_memory, float *v_memory, + long num_parameters, float learning_rate, + float beta1, float beta2, float beta1_correction, + float beta2_correction, float eps, float weight_decay) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= num_parameters) return; + float grad = grads_memory[i]; + float m = m_memory[i]; + float v = v_memory[i]; + m = lerp(grad, m, beta1); + m_memory[i] = m; + v = lerp(grad * grad, v, beta2); + v_memory[i] = v; + m /= beta1_correction; + v /= beta2_correction; + params_memory[i] -= learning_rate * (m / (sqrtf(v) + eps) + weight_decay * params_memory[i]); +} + +struct SoftmaxParams { + float Scale; + float Offset; +}; + +__device__ SoftmaxParams prepare_softmax_blockwide_nofloat4( + cg::thread_block_tile<32> &warp, int idx, const float *inp, int V, int P) { + const float *x = inp + idx * P; + float thread_maxval = -INFINITY; + float thread_sumval = 0.0f; + for (int i = V + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) { + float v = x[i]; + float old_maxval = thread_maxval; + thread_maxval = fmaxf(thread_maxval, v); + thread_sumval *= expf((old_maxval - thread_maxval)); + thread_sumval += expf(v - thread_maxval); + } + + __shared__ float shared_maxval[32]; + __shared__ float shared_sumval[32]; + int num_warps = blockDim.x / 32; + int warp_id = threadIdx.x / 32; + int lane_id = threadIdx.x % 32; + + float warp_maxval = cg::reduce(warp, thread_maxval, cg::greater{}); + if (lane_id == 0) shared_maxval[warp_id] = warp_maxval; + __syncthreads(); + warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX; + float block_maxval = cg::reduce(warp, warp_maxval, cg::greater{}); + thread_sumval *= expf(thread_maxval - block_maxval); + float warp_sumval = cg::reduce(warp, thread_sumval, cg::plus{}); + if (lane_id == 0) shared_sumval[warp_id] = warp_sumval; + __syncthreads(); + warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f; + float block_sumval = cg::reduce(warp, warp_sumval, cg::plus{}); + return SoftmaxParams{1.f / block_sumval, block_maxval}; +} + +__global__ void fused_classifier_kernel3(float *logits, float *losses, + float *probs, const float *dlosses, + const int *targets, int B, int T, + int V, int P) { + cg::thread_block block = cg::this_thread_block(); + cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); + int idx = blockIdx.x; + int ix = targets[idx]; + + SoftmaxParams sp = prepare_softmax_blockwide_nofloat4(warp, idx, logits, V, P); + + if (threadIdx.x == 0) { + float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale; + losses[idx] = -logf(prob); + } + + float dloss = dlosses != NULL ? dlosses[idx] : 1.0f / (B * T); + const float *logits_vec = logits + idx * P; + for (int i = threadIdx.x; i < V; i += blockDim.x) { + float v = __ldcs(&logits_vec[i]); + float prob = expf(v - sp.Offset) * sp.Scale; + if (probs != NULL) { + probs[idx * P + i] = prob; + } + float indicator = (i == ix) ? 1.0f : 0.0f; + logits[idx * P + i] = (prob - indicator) * dloss; + } +} + +__device__ float4 ld_vec(const float *address) { + return *reinterpret_cast(address); +} + +__device__ void st_vec(float *address, float4 val) { + *reinterpret_cast(address) = val; +} + +__global__ void __launch_bounds__(16 * 16, 2) + matmul_forward_kernel4(float *out, const float *inp, const float *weight, + const float *bias, int C, int OC) { + int oc = 8 * (blockIdx.y * blockDim.y + threadIdx.y); + + __shared__ float lhs_s[128][32]; + __shared__ float rhs_s[128][32]; + + inp += 128 * blockIdx.x * C; + weight += 128 * blockIdx.y * C; + out += 128 * blockIdx.x * OC + 128 * blockIdx.y; + + float vals[8][8] = {}; + if (bias != NULL) { + for (int i = 0; i < 8; i++) { + for (int j = 0; j < 8; j += 4) { + float4 b = ld_vec(bias + oc + j); + vals[i][j + 0] = b.x; + vals[i][j + 1] = b.y; + vals[i][j + 2] = b.z; + vals[i][j + 3] = b.w; + } + } + } + + int si_start = 4 * (16 * threadIdx.y + threadIdx.x); + for (int so = 0; so < C; so += 32) { + __syncthreads(); + int xmod8 = threadIdx.x % 8; + int xby8 = threadIdx.x / 8; + int xo = 4 * xmod8; + for (int y = 2 * threadIdx.y + xby8; y < 128; y += 32) { + st_vec(&lhs_s[y][xo], ld_vec(inp + y * C + so + xo)); + st_vec(&rhs_s[y][xo], ld_vec(weight + y * C + so + xo)); + } + __syncthreads(); + + for (int si = si_start; si < si_start + 32; si += 4) { + float4 rhs[8]; + for (int u = 0; u < 8; ++u) { + rhs[u] = ld_vec(&rhs_s[u + 8 * threadIdx.y][si % 32]); + } + for (int ii = 0; ii < 8; ++ii) { + float4 lhs = ld_vec(&lhs_s[ii + 8 * threadIdx.x][si % 32]); + for (int ji = 0; ji < 8; ++ji) { + vals[ii][ji] += lhs.x * rhs[ji].x; + vals[ii][ji] += lhs.y * rhs[ji].y; + vals[ii][ji] += lhs.z * rhs[ji].z; + vals[ii][ji] += lhs.w * rhs[ji].w; + } + } + } + } + + for (int i = 0; i < 8; ++i) { + for (int j = 0; j < 8; j += 4) { + float4 result; + result.x = vals[i][j + 0]; + result.y = vals[i][j + 1]; + result.z = vals[i][j + 2]; + result.w = vals[i][j + 3]; + st_vec(out + (8 * threadIdx.x + i) * OC + 8 * threadIdx.y + j, result); + } + } +} + +// ---------------------------------------------------------------------------- +// kernel launchers + +void encoder_forward(float *out, const int *inp, const float *wte, + const float *wpe, int B, int T, int C) { + assert(C % 4 == 0); + const int block_size = 512; + const int N = B * T * C; + const int grid_size = CEIL_DIV(N / 4, block_size); + encoder_forward_kernel3<<>>( + (float4 *)out, inp, (float4 *)wte, (float4 *)wpe, B, T, C); + cudaCheck(cudaGetLastError()); +} + +void encoder_backward(float *dwte, float *dwpe, const float *dout, + const int *inp, int B, int T, int C) { + const int N = B * T * C; + const int block_size = 256; + const int grid_size = CEIL_DIV(N, block_size); + encoder_backward_kernel<<>>(dwte, dwpe, dout, inp, B, T, C); + cudaCheck(cudaGetLastError()); +} + +void layernorm_forward(float *out, float *mean, float *rstd, float *inp, + float *weight, float *bias, int B, int T, int C) { + const int block_size = 512; + const int N = B * T; + const int grid_size = CEIL_DIV(N * 32, block_size); + layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); + cudaCheck(cudaGetLastError()); +} + +void matmul_forward(float *out, const float *inp, const float *weight, + const float *bias, int B, int T, int C, int OC) { + int sqrt_block_size = 16; + dim3 gridDim(CEIL_DIV(B * T, 8 * sqrt_block_size), CEIL_DIV(OC, 8 * sqrt_block_size)); + dim3 blockDim(sqrt_block_size, sqrt_block_size); + matmul_forward_kernel4<<>>(out, inp, weight, bias, C, OC); + cudaCheck(cudaGetLastError()); +} + +void attention_forward(float *out, float *qkvr, float *att, float *inp, int B, + int T, int C, int NH) { + const int block_size = 256; + const int softmax_block_size = 256; + int HS = C / NH; + + float *q, *k, *v; + q = qkvr + 0 * B * T * C; + k = qkvr + 1 * B * T * C; + v = qkvr + 2 * B * T * C; + int total_threads = B * NH * T * HS; + int num_blocks = CEIL_DIV(total_threads, block_size); + permute_kernel<<>>(q, k, v, inp, B, T, NH, HS); + cudaCheck(cudaGetLastError()); + + const float alpha = 1.0f; + const float beta = 0.0f; + float *preatt = inp; + cublasCheck(cublasSgemmStridedBatched( + cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, k, HS, T * HS, + q, HS, T * HS, &beta, preatt, T, T * T, B * NH)); + + float scale = 1.0 / sqrtf(HS); + int grid_size = CEIL_DIV(B * NH * T * 32, softmax_block_size); + softmax_forward_kernel5<<>>(att, scale, preatt, B * NH, T); + cudaCheck(cudaGetLastError()); + + float *vaccum = inp; + cublasCheck(cublasSgemmStridedBatched( + cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha, v, HS, T * HS, + att, T, T * T, &beta, vaccum, HS, T * HS, B * NH)); + + num_blocks = CEIL_DIV(B * T * C, block_size); + unpermute_kernel<<>>(vaccum, out, B, T, NH, HS); + cudaCheck(cudaGetLastError()); +} + +void residual_forward(float *out, float *inp1, float *inp2, int N) { + const int block_size = 256; + const int grid_size = CEIL_DIV(N, block_size); + residual_forward_kernel<<>>(out, inp1, inp2, N); + cudaCheck(cudaGetLastError()); +} + +void gelu_forward(float *out, const float *inp, int N) { + const int block_size = 128; + const int grid_size = CEIL_DIV(N, block_size); + gelu_forward_kernel<<>>(out, inp, N); + cudaCheck(cudaGetLastError()); +} + +void gelu_backward(float *dinp, const float *inp, const float *dout, const int N) { + const int block_size = 128; + const int grid_size = CEIL_DIV(N, block_size); + gelu_backward_kernel<<>>(dinp, inp, dout, N); + cudaCheck(cudaGetLastError()); +} + +void matmul_backward(float *dinp, float *dweight, float *dbias, float *dout, + float *inp, float *weight, int B, int T, int C, int OC) { + float one = 1.0f; + float zero = 0.0f; + cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B * T, OC, + &one, weight, C, dout, OC, &zero, dinp, C)); + cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B * T, + &one, inp, C, dout, OC, &one, dweight, C)); + if (dbias != NULL) { + const int block_size = 1024; + const int grid_size = OC / 32; + matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); + cudaCheck(cudaGetLastError()); + } +} + +void layernorm_backward(float *dinp, float *dweight, float *dbias, + const float *dout, const float *inp, + const float *weight, const float *mean, + const float *rstd, int B, int T, int C) { + const int block_size = 512; + const int N = B * T; + const int grid_size = CEIL_DIV(32 * N, block_size); + size_t shared_mem_size = 2 * C * sizeof(float); + layernorm_backward_kernel2<<>>( + dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); + cudaCheck(cudaGetLastError()); +} + +void attention_backward(float *dinp, float *dqkvr, float *dpreatt, float *datt, + float *scratch, const float *dout, const float *qkvr, + const float *att, int B, int T, int C, int NH) { + const int block_size = 256; + int HS = C / NH; + const float one = 1.0f; + const float zero = 0.0f; + const float *q, *k, *v; + q = qkvr + 0 * B * T * C; + k = qkvr + 1 * B * T * C; + v = qkvr + 2 * B * T * C; + float *dq, *dk, *dv; + dq = dqkvr + 0 * B * T * C; + dk = dqkvr + 1 * B * T * C; + dv = dqkvr + 2 * B * T * C; + int num_blocks = CEIL_DIV(B * T * C, block_size); + unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); + cudaCheck(cudaGetLastError()); + cublasCheck(cublasSgemmStridedBatched( + cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &one, v, HS, T * HS, + scratch, HS, T * HS, &zero, datt, T, T * T, B * NH)); + cublasCheck(cublasSgemmStridedBatched( + cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, scratch, HS, + T * HS, att, T, T * T, &zero, dv, HS, T * HS, B * NH)); + int hs = C / NH; + float scale = 1.0f / sqrtf(hs); + softmax_autoregressive_backward_kernel<<>>( + dpreatt, datt, att, B, T, C, scale); + cudaCheck(cudaGetLastError()); + cublasCheck(cublasSgemmStridedBatched( + cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &one, k, HS, T * HS, + dpreatt, T, T * T, &zero, dq, HS, T * HS, B * NH)); + cublasCheck(cublasSgemmStridedBatched( + cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, q, HS, T * HS, + dpreatt, T, T * T, &zero, dk, HS, T * HS, B * NH)); + num_blocks = CEIL_DIV(B * NH * T * HS, block_size); + permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); + cudaCheck(cudaGetLastError()); +} + +void fused_classifier3(float *logits, float *losses, const float *dlosses, + const int *targets, int B, int T, int V, int P) { + const int block_size = 1024; + const int N = B * T; + const int grid_size = N; + fused_classifier_kernel3<<>>(logits, losses, NULL, dlosses, targets, B, T, V, P); + cudaCheck(cudaGetLastError()); +} + +// ---------------------------------------------------------------------------- +// GPT-2 model definition + +typedef struct { + int max_seq_len; + int vocab_size; + int padded_vocab_size; + int num_layers; + int num_heads; + int channels; +} GPT2Config; + +#define NUM_PARAMETER_TENSORS 16 +typedef struct { + float *wte; // (V, C) + float *wpe; // (maxT, C) - only on PE 0 + float *ln1w; // (L_local, C) + float *ln1b; // (L_local, C) + float *qkvw; // (L_local, 3*C, C) + float *qkvb; // (L_local, 3*C) + float *attprojw; // (L_local, C, C) + float *attprojb; // (L_local, C) + float *ln2w; // (L_local, C) + float *ln2b; // (L_local, C) + float *fcw; // (L_local, 4*C, C) + float *fcb; // (L_local, 4*C) + float *fcprojw; // (L_local, C, 4*C) + float *fcprojb; // (L_local, C) + float *lnfw; // (C) - only on final PE + float *lnfb; // (C) - only on final PE +} ParameterTensors; + +// Full parameter sizes (for file offset calculation) +void fill_in_parameter_sizes(size_t *param_sizes, GPT2Config config) { + int Vp = config.padded_vocab_size; + int C = config.channels; + int maxT = config.max_seq_len; + int L = config.num_layers; + param_sizes[0] = Vp * C; + param_sizes[1] = maxT * C; + param_sizes[2] = L * C; + param_sizes[3] = L * C; + param_sizes[4] = L * (3 * C) * C; + param_sizes[5] = L * (3 * C); + param_sizes[6] = L * C * C; + param_sizes[7] = L * C; + param_sizes[8] = L * C; + param_sizes[9] = L * C; + param_sizes[10] = L * (4 * C) * C; + param_sizes[11] = L * (4 * C); + param_sizes[12] = L * C * (4 * C); + param_sizes[13] = L * C; + param_sizes[14] = C; + param_sizes[15] = C; +} + +// Partitioned parameter sizes (for allocation) +void fill_in_parameter_sizes_partitioned(size_t *param_sizes, GPT2Config config, + PipelinePartition *part) { + int Vp = config.padded_vocab_size; + int C = config.channels; + int maxT = config.max_seq_len; + int L = part->num_layers; + + param_sizes[0] = Vp * C; // wte - all PEs + param_sizes[1] = part->has_embedding ? maxT * C : 0; // wpe + param_sizes[2] = L * C; + param_sizes[3] = L * C; + param_sizes[4] = L * (3 * C) * C; + param_sizes[5] = L * (3 * C); + param_sizes[6] = L * C * C; + param_sizes[7] = L * C; + param_sizes[8] = L * C; + param_sizes[9] = L * C; + param_sizes[10] = L * (4 * C) * C; + param_sizes[11] = L * (4 * C); + param_sizes[12] = L * C * (4 * C); + param_sizes[13] = L * C; + param_sizes[14] = part->has_final_ln ? C : 0; + param_sizes[15] = part->has_final_ln ? C : 0; +} + +float *malloc_and_point_parameters(ParameterTensors *params, + size_t *param_sizes, int on_device) { + size_t num_parameters = 0; + for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { + num_parameters += param_sizes[i]; + } + float *params_memory; + if (on_device) { + cudaCheck(cudaMalloc((void **)¶ms_memory, num_parameters * sizeof(float))); + } else { + params_memory = (float *)mallocCheck(num_parameters * sizeof(float)); + } + float **ptrs[] = { + ¶ms->wte, ¶ms->wpe, ¶ms->ln1w, ¶ms->ln1b, + ¶ms->qkvw, ¶ms->qkvb, ¶ms->attprojw, ¶ms->attprojb, + ¶ms->ln2w, ¶ms->ln2b, ¶ms->fcw, ¶ms->fcb, + ¶ms->fcprojw, ¶ms->fcprojb, ¶ms->lnfw, ¶ms->lnfb}; + float *params_memory_iterator = params_memory; + for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { + *(ptrs[i]) = params_memory_iterator; + params_memory_iterator += param_sizes[i]; + } + return params_memory; +} + +#define NUM_ACTIVATION_TENSORS 22 +typedef struct { + float *encoded; // (B, T, C) - only PE 0 + float *ln1; // (L_local, B, T, C) + float *ln1_mean; // (L_local, B, T) + float *ln1_rstd; // (L_local, B, T) + float *atty; // (L_local, B, T, C) + float *att; // (L_local, B, NH, T, T) + float *attproj; // (L_local, B, T, C) + float *residual2; // (L_local, B, T, C) + float *ln2; // (L_local, B, T, C) + float *ln2_mean; // (L_local, B, T) + float *ln2_rstd; // (L_local, B, T) + float *fch; // (L_local, B, T, 4*C) + float *fch_gelu; // (L_local, B, T, 4*C) + float *fcproj; // (L_local, B, T, C) + float *residual3; // (L_local, B, T, C) + float *lnf; // (B, T, C) - only final PE + float *lnf_mean; // (B, T) - only final PE + float *lnf_rstd; // (B, T) - only final PE + float *losses; // (B, T) - only final PE + float *qkvr; // (L_local, B, T, 3*C) + float *output; // scratch buffer + float *scratch_btc; // (B, T, C) - extra scratch for backward on PE 0 +} ActivationTensors; + +void fill_in_activation_sizes_partitioned(size_t *act_sizes, int B, int T, + GPT2Config config, PipelinePartition *part) { + size_t Vp = config.padded_vocab_size; + size_t L = part->num_layers; + size_t NH = config.num_heads; + size_t C = config.channels; + + act_sizes[0] = part->has_embedding ? B * T * C : 0; + act_sizes[1] = L * B * T * C; + act_sizes[2] = L * B * T; + act_sizes[3] = L * B * T; + act_sizes[4] = L * B * T * C; + act_sizes[5] = L * B * NH * T * T; + act_sizes[6] = L * B * T * C; + act_sizes[7] = L * B * T * C; + act_sizes[8] = L * B * T * C; + act_sizes[9] = L * B * T; + act_sizes[10] = L * B * T; + act_sizes[11] = L * B * T * 4 * C; + act_sizes[12] = L * B * T * 4 * C; + act_sizes[13] = L * B * T * C; + act_sizes[14] = L * B * T * C; + act_sizes[15] = part->has_final_ln ? B * T * C : 0; + act_sizes[16] = part->has_final_ln ? B * T : 0; + act_sizes[17] = part->has_final_ln ? B * T : 0; + act_sizes[18] = part->has_final_ln ? B * T : 0; + act_sizes[19] = L * B * T * 3 * C; + act_sizes[20] = B * T * max(3 * C, max(NH * T, Vp)); + act_sizes[21] = B * T * C; // scratch_btc for backward +} + +#define NUM_BACKWARD_TENSORS 3 +typedef struct { + float *bt4c; + float *preatt; + float *residual3; +} GradActTensors; + +void fill_in_grad_act_sizes(size_t *act_sizes, int B, int T, GPT2Config config) { + size_t NH = config.num_heads; + size_t C = config.channels; + act_sizes[0] = B * T * 4 * C; + act_sizes[1] = B * NH * T * T; + act_sizes[2] = B * T * C; +} + +float *malloc_and_point(float **targets[], const size_t *act_sizes, int n) { + size_t num_activations = 0; + for (size_t i = 0; i < n; i++) { + num_activations += act_sizes[i]; + } + float *acts_memory; + cudaCheck(cudaMalloc((void **)&acts_memory, num_activations * sizeof(float))); + float *acts_memory_iterator = acts_memory; + for (size_t i = 0; i < n; i++) { + *(targets[i]) = acts_memory_iterator; + acts_memory_iterator += act_sizes[i]; + } + return acts_memory; +} + +float *malloc_and_point_activations(ActivationTensors *acts, const size_t *act_sizes) { + float **ptrs[] = {&acts->encoded, &acts->ln1, &acts->ln1_mean, + &acts->ln1_rstd, &acts->atty, &acts->att, + &acts->attproj, &acts->residual2, &acts->ln2, + &acts->ln2_mean, &acts->ln2_rstd, &acts->fch, + &acts->fch_gelu, &acts->fcproj, &acts->residual3, + &acts->lnf, &acts->lnf_mean, &acts->lnf_rstd, + &acts->losses, &acts->qkvr, &acts->output, + &acts->scratch_btc}; + return malloc_and_point(ptrs, act_sizes, NUM_ACTIVATION_TENSORS); +} + +float *malloc_and_point_backward(GradActTensors *acts, const size_t *act_sizes) { + float **ptrs[] = {&acts->bt4c, &acts->preatt, &acts->residual3}; + return malloc_and_point(ptrs, act_sizes, NUM_BACKWARD_TENSORS); +} + +typedef struct { + GPT2Config config; + PipelinePartition partition; + ParameterTensors params; + size_t param_sizes[NUM_PARAMETER_TENSORS]; + float *params_memory; + size_t num_parameters; + ParameterTensors grads; + float *grads_memory; + float *m_memory; + float *v_memory; + ActivationTensors acts; + size_t act_sizes[NUM_ACTIVATION_TENSORS]; + float *acts_memory; + size_t num_activations; + GradActTensors grads_acts; + size_t num_grad_acts; + float *grads_acts_memory; + int batch_size; + int seq_len; + int *inputs; + int *targets; + float mean_loss; + float *cpu_losses; + float *nvshmem_act_buffer; + float *nvshmem_grad_buffer; + size_t nvshmem_buffer_size; + ncclComm_t nccl_comm; +} GPT2; + +void gpt2_build_from_checkpoint_partitioned(GPT2 *model, const char *checkpoint_path, + PipelinePartition *part) { + int my_pe = nvshmem_my_pe(); + + FILE *model_file = fopenCheck(checkpoint_path, "rb"); + int model_header[256]; + freadCheck(model_header, sizeof(int), 256, model_file); + if (model_header[0] != 20240326) { + fprintf(stderr, "Bad magic model file\n"); + exit(EXIT_FAILURE); + } + if (model_header[1] != 3) { + fprintf(stderr, "Bad version in model file\n"); + exit(EXIT_FAILURE); + } + + 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]; + model->config.padded_vocab_size = model_header[7]; + model->partition = *part; + + int C = model->config.channels; + int Vp = model->config.padded_vocab_size; + int maxT = model->config.max_seq_len; + int L_total = model->config.num_layers; + + // Full sizes for file offset calculation + size_t full_param_sizes[NUM_PARAMETER_TENSORS]; + fill_in_parameter_sizes(full_param_sizes, model->config); + + // Partitioned sizes for allocation + fill_in_parameter_sizes_partitioned(model->param_sizes, model->config, part); + + size_t num_parameters = 0; + for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { + num_parameters += model->param_sizes[i]; + } + model->num_parameters = num_parameters; + + model->params_memory = malloc_and_point_parameters(&model->params, model->param_sizes, 1); + + printf("[PE %d] allocated %zu MiB for partitioned parameters (layers %d-%d)\n", + my_pe, (num_parameters * sizeof(float)) >> 20, + part->first_layer, part->first_layer + part->num_layers - 1); + + size_t header_size = 256 * sizeof(int); + + // Helper to read a range of floats from file + auto read_range = [&](float *gpu_dst, size_t file_elem_offset, size_t count) { + if (count == 0) return; + float *cpu_buf = (float *)mallocCheck(count * sizeof(float)); + fseekCheck(model_file, header_size + file_elem_offset * sizeof(float), SEEK_SET); + freadCheck(cpu_buf, sizeof(float), count, model_file); + cudaCheck(cudaMemcpy(gpu_dst, cpu_buf, count * sizeof(float), cudaMemcpyHostToDevice)); + free(cpu_buf); + }; + + size_t file_offset = 0; + size_t layer_start = part->first_layer; + size_t layer_count = part->num_layers; + + // 0: wte - all PEs need this + read_range(model->params.wte, file_offset, Vp * C); + file_offset += full_param_sizes[0]; + + // 1: wpe - only PE 0 + if (part->has_embedding) { + read_range(model->params.wpe, file_offset, maxT * C); + } + file_offset += full_param_sizes[1]; + + // 2: ln1w + read_range(model->params.ln1w, file_offset + layer_start * C, layer_count * C); + file_offset += full_param_sizes[2]; + + // 3: ln1b + read_range(model->params.ln1b, file_offset + layer_start * C, layer_count * C); + file_offset += full_param_sizes[3]; + + // 4: qkvw + read_range(model->params.qkvw, file_offset + layer_start * 3 * C * C, layer_count * 3 * C * C); + file_offset += full_param_sizes[4]; + + // 5: qkvb + read_range(model->params.qkvb, file_offset + layer_start * 3 * C, layer_count * 3 * C); + file_offset += full_param_sizes[5]; + + // 6: attprojw + read_range(model->params.attprojw, file_offset + layer_start * C * C, layer_count * C * C); + file_offset += full_param_sizes[6]; + + // 7: attprojb + read_range(model->params.attprojb, file_offset + layer_start * C, layer_count * C); + file_offset += full_param_sizes[7]; + + // 8: ln2w + read_range(model->params.ln2w, file_offset + layer_start * C, layer_count * C); + file_offset += full_param_sizes[8]; + + // 9: ln2b + read_range(model->params.ln2b, file_offset + layer_start * C, layer_count * C); + file_offset += full_param_sizes[9]; + + // 10: fcw + read_range(model->params.fcw, file_offset + layer_start * 4 * C * C, layer_count * 4 * C * C); + file_offset += full_param_sizes[10]; + + // 11: fcb + read_range(model->params.fcb, file_offset + layer_start * 4 * C, layer_count * 4 * C); + file_offset += full_param_sizes[11]; + + // 12: fcprojw + read_range(model->params.fcprojw, file_offset + layer_start * C * 4 * C, layer_count * C * 4 * C); + file_offset += full_param_sizes[12]; + + // 13: fcprojb + read_range(model->params.fcprojb, file_offset + layer_start * C, layer_count * C); + file_offset += full_param_sizes[13]; + + // 14: lnfw - only final PE + if (part->has_final_ln) { + read_range(model->params.lnfw, file_offset, C); + } + file_offset += full_param_sizes[14]; + + // 15: lnfb - only final PE + if (part->has_final_ln) { + read_range(model->params.lnfb, file_offset, C); + } + + fcloseCheck(model_file); + + model->acts_memory = NULL; + model->grads_memory = NULL; + model->m_memory = NULL; + model->v_memory = NULL; + model->grads_acts_memory = NULL; + model->inputs = NULL; + model->targets = NULL; + model->cpu_losses = NULL; + model->batch_size = 0; + model->seq_len = 0; + model->mean_loss = -1.0f; + model->nvshmem_act_buffer = NULL; + model->nvshmem_grad_buffer = NULL; + model->nvshmem_buffer_size = 0; +} + +void gpt2_allocate_nvshmem_buffers(GPT2 *model) { + int my_pe = nvshmem_my_pe(); + int B = model->batch_size; + int T = model->seq_len; + int C = model->config.channels; + + size_t buffer_elements = B * T * C; + model->nvshmem_buffer_size = buffer_elements * sizeof(float); + + model->nvshmem_act_buffer = (float *)nvshmem_malloc(model->nvshmem_buffer_size); + model->nvshmem_grad_buffer = (float *)nvshmem_malloc(model->nvshmem_buffer_size); + + if (model->nvshmem_act_buffer == NULL || model->nvshmem_grad_buffer == NULL) { + printf("[PE %d] Error: Failed to allocate NVSHMEM buffers\n", my_pe); + exit(EXIT_FAILURE); + } + + printf("[PE %d] allocated %zu MiB for NVSHMEM buffers\n", my_pe, + (2 * model->nvshmem_buffer_size) >> 20); + + cudaCheck(cudaMemset(model->nvshmem_act_buffer, 0, model->nvshmem_buffer_size)); + cudaCheck(cudaMemset(model->nvshmem_grad_buffer, 0, model->nvshmem_buffer_size)); +} + +void gpt2_forward(GPT2 *model, int *inputs, int *targets, int B, int T) { + if (model->params_memory == NULL) { + printf("Error: model was not initialized properly.\n"); + exit(EXIT_FAILURE); + } + + int V = model->config.vocab_size; + int Vp = model->config.padded_vocab_size; + int NH = model->config.num_heads; + int C = model->config.channels; + PipelinePartition part = model->partition; + int my_pe = nvshmem_my_pe(); + + for (int i = 0; i < B * T; i++) { + assert(0 <= inputs[i] && inputs[i] < V); + if (targets != NULL) { + assert(0 <= targets[i] && targets[i] < V); + } + } + + if (model->acts_memory == NULL) { + model->batch_size = B; + model->seq_len = T; + fill_in_activation_sizes_partitioned(model->act_sizes, B, T, model->config, &part); + size_t num_activations = 0; + for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) { + num_activations += model->act_sizes[i]; + } + model->num_activations = num_activations; + model->acts_memory = malloc_and_point_activations(&model->acts, model->act_sizes); + printf("[PE %d] allocated %zu MiB for activations\n", my_pe, + (num_activations * sizeof(float)) >> 20); + cudaCheck(cudaMalloc((void **)&model->inputs, B * T * sizeof(int))); + cudaCheck(cudaMalloc((void **)&model->targets, B * T * sizeof(int))); + cudaCheck(cudaMallocHost((void **)&model->cpu_losses, B * T * sizeof(float))); + } else { + if (B != model->batch_size || T != model->seq_len) { + printf("Model: B=%d T=%d, Desired: B=%d T=%d\n", model->batch_size, model->seq_len, B, T); + exit(EXIT_FAILURE); + } + } + + cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice)); + if (targets != NULL) { + cudaCheck(cudaMemcpy(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice)); + } + + ParameterTensors params = model->params; + ActivationTensors acts = model->acts; + float *residual; + + // GPU 0: Embedding + first half of layers + if (my_pe == 0) { + encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); + + for (int l = 0; l < part.num_layers; l++) { + residual = (l == 0) ? acts.encoded : acts.residual3 + (l - 1) * B * T * C; + + float *l_ln1w = params.ln1w + l * C; + float *l_ln1b = params.ln1b + l * C; + float *l_qkvw = params.qkvw + l * 3 * C * C; + float *l_qkvb = params.qkvb + l * 3 * C; + float *l_attprojw = params.attprojw + l * C * C; + float *l_attprojb = params.attprojb + l * C; + float *l_ln2w = params.ln2w + l * C; + float *l_ln2b = params.ln2b + l * C; + float *l_fcw = params.fcw + l * 4 * C * C; + float *l_fcb = params.fcb + l * 4 * C; + float *l_fcprojw = params.fcprojw + l * C * 4 * C; + float *l_fcprojb = params.fcprojb + l * C; + + float *l_ln1 = acts.ln1 + l * B * T * C; + float *l_ln1_mean = acts.ln1_mean + l * B * T; + float *l_ln1_rstd = acts.ln1_rstd + l * B * T; + float *l_qkvr = acts.qkvr + l * B * T * 3 * C; + float *l_atty = acts.atty + l * B * T * C; + float *l_att = acts.att + l * B * NH * T * T; + float *l_attproj = acts.attproj + l * B * T * C; + float *l_residual2 = acts.residual2 + l * B * T * C; + float *l_ln2 = acts.ln2 + l * B * T * C; + float *l_ln2_mean = acts.ln2_mean + l * B * T; + float *l_ln2_rstd = acts.ln2_rstd + l * B * T; + float *l_fch = acts.fch + l * B * T * 4 * C; + float *l_fch_gelu = acts.fch_gelu + l * B * T * 4 * C; + float *l_fcproj = acts.fcproj + l * B * T * C; + float *l_residual3 = acts.residual3 + l * B * T * C; + float *scratch = acts.output; + + layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C); + matmul_forward(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3 * C); + attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH); + matmul_forward(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C); + residual_forward(l_residual2, residual, l_attproj, B * T * C); + layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C); + matmul_forward(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4 * C); + gelu_forward(l_fch_gelu, l_fch, B * T * 4 * C); + matmul_forward(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4 * C, C); + residual_forward(l_residual3, l_residual2, l_fcproj, B * T * C); + } + + float *last_output = acts.residual3 + (part.num_layers - 1) * B * T * C; + cudaCheck(cudaDeviceSynchronize()); + nvshmem_putmem(model->nvshmem_act_buffer, last_output, B * T * C * sizeof(float), 1); + nvshmem_quiet(); + + if (targets != NULL) { + model->mean_loss = 0.0f; + } else { + model->mean_loss = -1.0f; + } + } + + nvshmem_barrier_all(); + + // GPU 1: Second half of layers + final LN + loss + if (my_pe == 1) { + for (int l = 0; l < part.num_layers; l++) { + if (l == 0) { + residual = model->nvshmem_act_buffer; + } else { + residual = acts.residual3 + (l - 1) * B * T * C; + } + + float *l_ln1w = params.ln1w + l * C; + float *l_ln1b = params.ln1b + l * C; + float *l_qkvw = params.qkvw + l * 3 * C * C; + float *l_qkvb = params.qkvb + l * 3 * C; + float *l_attprojw = params.attprojw + l * C * C; + float *l_attprojb = params.attprojb + l * C; + float *l_ln2w = params.ln2w + l * C; + float *l_ln2b = params.ln2b + l * C; + float *l_fcw = params.fcw + l * 4 * C * C; + float *l_fcb = params.fcb + l * 4 * C; + float *l_fcprojw = params.fcprojw + l * C * 4 * C; + float *l_fcprojb = params.fcprojb + l * C; + + float *l_ln1 = acts.ln1 + l * B * T * C; + float *l_ln1_mean = acts.ln1_mean + l * B * T; + float *l_ln1_rstd = acts.ln1_rstd + l * B * T; + float *l_qkvr = acts.qkvr + l * B * T * 3 * C; + float *l_atty = acts.atty + l * B * T * C; + float *l_att = acts.att + l * B * NH * T * T; + float *l_attproj = acts.attproj + l * B * T * C; + float *l_residual2 = acts.residual2 + l * B * T * C; + float *l_ln2 = acts.ln2 + l * B * T * C; + float *l_ln2_mean = acts.ln2_mean + l * B * T; + float *l_ln2_rstd = acts.ln2_rstd + l * B * T; + float *l_fch = acts.fch + l * B * T * 4 * C; + float *l_fch_gelu = acts.fch_gelu + l * B * T * 4 * C; + float *l_fcproj = acts.fcproj + l * B * T * C; + float *l_residual3 = acts.residual3 + l * B * T * C; + float *scratch = acts.output; + + layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C); + matmul_forward(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3 * C); + attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH); + matmul_forward(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C); + residual_forward(l_residual2, residual, l_attproj, B * T * C); + layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C); + matmul_forward(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4 * C); + gelu_forward(l_fch_gelu, l_fch, B * T * 4 * C); + matmul_forward(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4 * C, C); + residual_forward(l_residual3, l_residual2, l_fcproj, B * T * C); + } + + residual = acts.residual3 + (part.num_layers - 1) * B * T * C; + layernorm_forward(acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, + params.lnfw, params.lnfb, B, T, C); + matmul_forward(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp); + + if (targets != NULL) { + fused_classifier3(acts.output, acts.losses, NULL, model->targets, B, T, V, Vp); + cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(float), + cudaMemcpyDeviceToHost)); + float mean_loss = 0.0f; + for (int i = 0; i < B * T; i++) { + mean_loss += model->cpu_losses[i]; + } + mean_loss /= B * T; + model->mean_loss = mean_loss; + } else { + model->mean_loss = -1.0f; + } + } + + nvshmem_barrier_all(); +} + +void gpt2_zero_grad(GPT2 *model) { + if (model->grads_acts_memory != NULL) { + cudaCheck(cudaMemset(model->grads_acts_memory, 0, model->num_grad_acts * sizeof(float))); + } + if (model->grads_memory != NULL) { + cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(float))); + } +} + +void gpt2_backward(GPT2 *model) { + counter++; + + if (model->mean_loss == -1.0f) { + printf("Error: must forward with targets before backward\n"); + exit(EXIT_FAILURE); + } + + PipelinePartition part = model->partition; + int my_pe = nvshmem_my_pe(); + + if (model->grads_memory == NULL) { + model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_sizes, 1); + printf("[PE %d] allocated %zu MiB for parameter gradients\n", my_pe, + (model->num_parameters * sizeof(float)) >> 20); + + size_t bw_act_sizes[NUM_BACKWARD_TENSORS]; + GPT2Config cfg = model->config; + cfg.num_layers = 1; + fill_in_grad_act_sizes(bw_act_sizes, model->batch_size, model->seq_len, cfg); + model->grads_acts_memory = malloc_and_point_backward(&model->grads_acts, bw_act_sizes); + model->num_grad_acts = 0; + for (int i = 0; i < NUM_BACKWARD_TENSORS; i++) { + model->num_grad_acts += bw_act_sizes[i]; + } + printf("[PE %d] allocated %zu MiB for activation gradients\n", my_pe, + (model->num_grad_acts * sizeof(float)) >> 20); + gpt2_zero_grad(model); + } + + int B = model->batch_size; + int T = model->seq_len; + int Vp = model->config.padded_vocab_size; + int NH = model->config.num_heads; + int C = model->config.channels; + + ParameterTensors params = model->params; + ParameterTensors grads = model->grads; + ActivationTensors acts = model->acts; + GradActTensors grads_acts = model->grads_acts; + float *residual; + float *dresidual = grads_acts.residual3; + + // GPU 1: Backward through classifier + second half layers + if (my_pe == 1) { + matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, + params.wte, B, T, C, Vp); + residual = acts.residual3 + (part.num_layers - 1) * B * T * C; + layernorm_backward(dresidual, grads.lnfw, grads.lnfb, grads_acts.bt4c, + residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C); + + for (int l = part.num_layers - 1; l >= 0; l--) { + if (l == 0) { + residual = model->nvshmem_act_buffer; + } else { + residual = acts.residual3 + (l - 1) * B * T * C; + } + + float *l_ln1w = params.ln1w + l * C; + float *l_qkvw = params.qkvw + l * 3 * C * C; + float *l_attprojw = params.attprojw + l * C * C; + float *l_ln2w = params.ln2w + l * C; + float *l_fcw = params.fcw + l * 4 * C * C; + float *l_fcprojw = params.fcprojw + l * C * 4 * C; + + float *dl_ln1w = grads.ln1w + l * C; + float *dl_ln1b = grads.ln1b + l * C; + float *dl_qkvw = grads.qkvw + l * 3 * C * C; + float *dl_qkvb = grads.qkvb + l * 3 * C; + float *dl_attprojw = grads.attprojw + l * C * C; + float *dl_attprojb = grads.attprojb + l * C; + float *dl_ln2w = grads.ln2w + l * C; + float *dl_ln2b = grads.ln2b + l * C; + float *dl_fcw = grads.fcw + l * 4 * C * C; + float *dl_fcb = grads.fcb + l * 4 * C; + float *dl_fcprojw = grads.fcprojw + l * C * 4 * C; + float *dl_fcprojb = grads.fcprojb + l * C; + + float *l_ln1 = acts.ln1 + l * B * T * C; + float *l_ln1_mean = acts.ln1_mean + l * B * T; + float *l_ln1_rstd = acts.ln1_rstd + l * B * T; + float *l_qkvr = acts.qkvr + l * B * T * 3 * C; + float *l_atty = acts.atty + l * B * T * C; + float *l_att = acts.att + l * B * NH * T * T; + float *l_residual2 = acts.residual2 + l * B * T * C; + float *l_ln2 = acts.ln2 + l * B * T * C; + float *l_ln2_mean = acts.ln2_mean + l * B * T; + float *l_ln2_rstd = acts.ln2_rstd + l * B * T; + float *l_fch = acts.fch + l * B * T * 4 * C; + float *l_fch_gelu = acts.fch_gelu + l * B * T * 4 * C; + + float *dl_btc = acts.lnf; + float *dl_bt4c = grads_acts.bt4c; + float *dl_preatt = grads_acts.preatt; + float *scratch = acts.output; + float *buffer_a = l_atty; + float *buffer_b = l_fch; + + matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, B, T, 4 * C, C); + gelu_backward(dl_bt4c, l_fch, dl_bt4c, B * T * 4 * C); + matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, B, T, C, 4 * C); + layernorm_backward(dresidual, dl_ln2w, dl_ln2b, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C); + matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, B, T, C, C); + attention_backward(dl_bt4c, buffer_b, dl_preatt, scratch, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH); + matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, B, T, C, 3 * C); + layernorm_backward(dresidual, dl_ln1w, dl_ln1b, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); + } + + cudaCheck(cudaDeviceSynchronize()); + nvshmem_putmem(model->nvshmem_grad_buffer, dresidual, B * T * C * sizeof(float), 0); + nvshmem_quiet(); + } + + nvshmem_barrier_all(); + + // GPU 0: Backward through first half layers + embedding + if (my_pe == 0) { + cudaMemcpy(dresidual, model->nvshmem_grad_buffer, B * T * C * sizeof(float), + cudaMemcpyDeviceToDevice); + + for (int l = part.num_layers - 1; l >= 0; l--) { + residual = (l == 0) ? acts.encoded : acts.residual3 + (l - 1) * B * T * C; + + float *l_ln1w = params.ln1w + l * C; + float *l_qkvw = params.qkvw + l * 3 * C * C; + float *l_attprojw = params.attprojw + l * C * C; + float *l_ln2w = params.ln2w + l * C; + float *l_fcw = params.fcw + l * 4 * C * C; + float *l_fcprojw = params.fcprojw + l * C * 4 * C; + + float *dl_ln1w = grads.ln1w + l * C; + float *dl_ln1b = grads.ln1b + l * C; + float *dl_qkvw = grads.qkvw + l * 3 * C * C; + float *dl_qkvb = grads.qkvb + l * 3 * C; + float *dl_attprojw = grads.attprojw + l * C * C; + float *dl_attprojb = grads.attprojb + l * C; + float *dl_ln2w = grads.ln2w + l * C; + float *dl_ln2b = grads.ln2b + l * C; + float *dl_fcw = grads.fcw + l * 4 * C * C; + float *dl_fcb = grads.fcb + l * 4 * C; + float *dl_fcprojw = grads.fcprojw + l * C * 4 * C; + float *dl_fcprojb = grads.fcprojb + l * C; + + float *l_ln1 = acts.ln1 + l * B * T * C; + float *l_ln1_mean = acts.ln1_mean + l * B * T; + float *l_ln1_rstd = acts.ln1_rstd + l * B * T; + float *l_qkvr = acts.qkvr + l * B * T * 3 * C; + float *l_atty = acts.atty + l * B * T * C; + float *l_att = acts.att + l * B * NH * T * T; + float *l_residual2 = acts.residual2 + l * B * T * C; + float *l_ln2 = acts.ln2 + l * B * T * C; + float *l_ln2_mean = acts.ln2_mean + l * B * T; + float *l_ln2_rstd = acts.ln2_rstd + l * B * T; + float *l_fch = acts.fch + l * B * T * 4 * C; + float *l_fch_gelu = acts.fch_gelu + l * B * T * 4 * C; + + // Use scratch_btc for dl_btc since GPU 0 doesn't have lnf + float *dl_btc = acts.scratch_btc; + float *dl_bt4c = grads_acts.bt4c; + float *dl_preatt = grads_acts.preatt; + float *scratch = acts.output; + float *buffer_a = l_atty; + float *buffer_b = l_fch; + + matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, B, T, 4 * C, C); + gelu_backward(dl_bt4c, l_fch, dl_bt4c, B * T * 4 * C); + matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, B, T, C, 4 * C); + layernorm_backward(dresidual, dl_ln2w, dl_ln2b, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C); + matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, B, T, C, C); + attention_backward(dl_bt4c, buffer_b, dl_preatt, scratch, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH); + matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, B, T, C, 3 * C); + layernorm_backward(dresidual, dl_ln1w, dl_ln1b, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C); + } + + encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C); + } + + nvshmem_barrier_all(); + + // All-reduce wte gradients + ncclCheck(ncclAllReduce((const void *)grads.wte, (void *)grads.wte, Vp * C, + ncclFloat, ncclSum, model->nccl_comm, 0)); + + nvshmem_barrier_all(); +} + +void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, + float eps, float weight_decay, int t) { + if (model->m_memory == NULL) { + cudaCheck(cudaMalloc((void **)&model->m_memory, model->num_parameters * sizeof(float))); + 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))); + int my_pe = nvshmem_my_pe(); + printf("[PE %d] allocated %zu MiB for AdamW optimizer state\n", my_pe, + (2 * model->num_parameters * sizeof(float)) >> 20); + } + + int block_size = 512; + int num_blocks = CEIL_DIV(model->num_parameters, block_size); + float beta1_correction = 1.0f - powf(beta1, t); + float beta2_correction = 1.0f - powf(beta2, t); + adamw_kernel2<<>>( + model->params_memory, model->grads_memory, model->m_memory, + model->v_memory, model->num_parameters, learning_rate, beta1, beta2, + beta1_correction, beta2_correction, eps, weight_decay); + cudaCheck(cudaGetLastError()); +} + +void gpt2_free(GPT2 *model) { + cudaCheck(cudaFree(model->params_memory)); + cudaCheck(cudaFree(model->grads_memory)); + cudaCheck(cudaFree(model->m_memory)); + cudaCheck(cudaFree(model->v_memory)); + cudaCheck(cudaFree(model->acts_memory)); + cudaCheck(cudaFree(model->grads_acts_memory)); + cudaCheck(cudaFree(model->inputs)); + cudaCheck(cudaFree(model->targets)); + cudaFreeHost(model->cpu_losses); + ncclCommDestroy(model->nccl_comm); +} + +#ifndef TESTING +// ---------------------------------------------------------------------------- +// sampler + +#define GPT2_EOT 50256 + +unsigned int random_u32(unsigned long long *state) { + *state ^= *state >> 12; + *state ^= *state << 25; + *state ^= *state >> 27; + return (*state * 0x2545F4914F6CDD1Dull) >> 32; +} + +float random_f32(unsigned long long *state) { + return (random_u32(state) >> 8) / 16777216.0f; +} + +int sample_softmax(const float *logits, int n, float coin) { + double norm = 0; + for (int i = 0; i < n; i++) { + norm += expf(logits[i]); + } + coin *= norm; + float cdf = 0.0f; + for (int i = 0; i < n; i++) { + cdf += expf(logits[i]); + if (coin < cdf) { + return i; + } + } + return n - 1; +} + +// ---------------------------------------------------------------------------- +// Logger + +typedef struct { + FILE *logfile; + int flush_every; +} Logger; + +void logger_init(Logger *logger, const char *filename) { + logger->flush_every = 20; + logger->logfile = NULL; + if (filename != NULL) { + logger->logfile = fopenCheck(filename, "w"); + } +} + +void logger_log_val(Logger *logger, int step, float val_loss) { + if (logger->logfile != NULL) { + fprintf(logger->logfile, "s:%d tel:%.4f\n", step, val_loss); + } +} + +void logger_log_train(Logger *logger, int step, float train_loss) { + if (logger->logfile != NULL) { + fprintf(logger->logfile, "s:%d trl:%.4f\n", step, train_loss); + if (step % 10 == 0) { + fflush(logger->logfile); + } + } +} + +void logger_free(Logger *logger) { + if (logger->logfile != NULL) { + fclose(logger->logfile); + } +} + +// ---------------------------------------------------------------------------- +// CLI + +void error_usage() { + fprintf(stderr, "Usage: ./train_gpt2_partitioned [options]\n"); + fprintf(stderr, "Options:\n"); + fprintf(stderr, " -i train data filename pattern\n"); + fprintf(stderr, " -j val data filename pattern\n"); + fprintf(stderr, " -o output log file (default = NULL)\n"); + fprintf(stderr, " -b batch size B (default = 4)\n"); + fprintf(stderr, " -t sequence length T (default = 1024)\n"); + fprintf(stderr, " -l learning rate (default = 3e-4f)\n"); + fprintf(stderr, " -v val_loss_every (default = 20)\n"); + fprintf(stderr, " -m val_max_steps (default = 20)\n"); + fprintf(stderr, " -s sample_every (default = 20)\n"); + fprintf(stderr, " -g genT (default = 64)\n"); + exit(EXIT_FAILURE); +} + +// ---------------------------------------------------------------------------- +// main training loop + +int main(int argc, char *argv[]) { + // Initialize MPI + MPI_Init(&argc, &argv); + + // Initialize NVSHMEM + nvshmemx_init_attr_t attr; + MPI_Comm mpi_comm = MPI_COMM_WORLD; + attr.mpi_comm = &mpi_comm; + nvshmemx_init_attr(NVSHMEMX_INIT_WITH_MPI_COMM, &attr); + + int my_pe = nvshmem_my_pe(); + int n_pes = nvshmem_n_pes(); + + if (n_pes != 2) { + if (my_pe == 0) { + fprintf(stderr, "Error: This pipeline requires exactly 2 GPUs, got %d\n", n_pes); + } + nvshmem_finalize(); + MPI_Finalize(); + return 1; + } + + cudaCheck(cudaSetDevice(my_pe)); + cudaDeviceProp deviceProp; + cudaGetDeviceProperties(&deviceProp, my_pe); + + cudaDeviceSynchronize(); + nvshmem_barrier_all(); + + // Initialize NCCL + ncclUniqueId nccl_id; + if (my_pe == 0) { + ncclCheck(ncclGetUniqueId(&nccl_id)); + } + MPI_Bcast(&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD); + ncclComm_t nccl_comm; + ncclCheck(ncclCommInitRank(&nccl_comm, n_pes, nccl_id, my_pe)); + + // Parse arguments + const char *train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; + const char *val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; + const char *output_log_file = NULL; + int B = 4; + int T = 1024; + float learning_rate = 3e-4f; + int val_loss_every = 20; + int val_max_steps = 20; + int sample_every = 20; + int genT = 64; + + for (int i = 1; i < argc; i += 2) { + if (i + 1 >= argc) error_usage(); + if (argv[i][0] != '-') error_usage(); + if (strlen(argv[i]) != 2) error_usage(); + if (argv[i][1] == 'i') train_data_pattern = argv[i + 1]; + else if (argv[i][1] == 'j') val_data_pattern = 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] == '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]); + else if (argv[i][1] == 'm') val_max_steps = atoi(argv[i + 1]); + else if (argv[i][1] == 's') sample_every = atoi(argv[i + 1]); + else if (argv[i][1] == 'g') genT = atoi(argv[i + 1]); + else error_usage(); + } + + if (my_pe == 0) { + printf("+------------------------+----------------------------------------------------+\n"); + printf("| Parameter | Value |\n"); + printf("+------------------------+----------------------------------------------------+\n"); + printf("| NVSHMEM PEs (GPUs) | %-50d |\n", n_pes); + printf("| train data pattern | %-50s |\n", train_data_pattern); + printf("| val data pattern | %-50s |\n", val_data_pattern); + printf("| batch size B | %-50d |\n", B); + printf("| sequence length T | %-50d |\n", T); + printf("| learning rate | %-50f |\n", learning_rate); + printf("+------------------------+----------------------------------------------------+\n"); + } + + // Setup cuBLAS + cublasCheck(cublasCreate(&cublas_handle)); + int enable_tf32 = deviceProp.major >= 8 ? 1 : 0; + cublas_compute_type = enable_tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F; + cublasMath_t cublas_math_mode = enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH; + cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode)); + + if (my_pe == 0) { + printf("| device | %-50s |\n", deviceProp.name); + printf("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); + printf("+------------------------+----------------------------------------------------+\n"); + } + + // Read header to get num_layers for partitioning + FILE *tmp = fopen("gpt2_124M.bin", "rb"); + if (tmp == NULL) { + printf("Error: could not open gpt2_124M.bin\n"); + exit(EXIT_FAILURE); + } + int header[256]; + fread(header, sizeof(int), 256, tmp); + fclose(tmp); + int total_layers = header[4]; + + // Compute partition for this PE + PipelinePartition part; + get_partition_for_pe(&part, my_pe, n_pes, total_layers); + + if (my_pe == 0) { + printf("| total_layers | %-50d |\n", total_layers); + printf("| layers per PE | %-50d |\n", part.num_layers); + printf("+------------------------+----------------------------------------------------+\n"); + } + + // Build model with partitioning + GPT2 model; + gpt2_build_from_checkpoint_partitioned(&model, "gpt2_124M.bin", &part); + model.nccl_comm = nccl_comm; + + if (my_pe == 0) { + printf("| max_sequence_length T | %-50d |\n", model.config.max_seq_len); + printf("| vocab_size V | %-50d |\n", model.config.vocab_size); + printf("| padded_vocab_size Vp | %-50d |\n", model.config.padded_vocab_size); + printf("| num_layers L (total) | %-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("+------------------------+----------------------------------------------------+\n"); + } + + // Build data loaders + DataLoader train_loader, val_loader; + dataloader_init(&train_loader, train_data_pattern, B, T, 0, 1, 1); + dataloader_init(&val_loader, val_data_pattern, B, T, 0, 1, 0); + int train_num_batches = train_loader.num_tokens / (B * T); + int val_num_batches = val_loader.num_tokens / (B * T); + if (val_num_batches > val_max_steps) val_num_batches = val_max_steps; + + if (my_pe == 0) { + printf("| train_num_batches | %-50d |\n", train_num_batches); + printf("| val_num_batches | %-50d |\n", val_num_batches); + printf("+------------------------+----------------------------------------------------+\n"); + } + + // Allocate NVSHMEM buffers + model.batch_size = B; + model.seq_len = T; + gpt2_allocate_nvshmem_buffers(&model); + + // Dummy forward to allocate activations + int *dummy_batch = (int *)malloc(B * T * sizeof(int)); + for (int i = 0; i < B * T; i++) dummy_batch[i] = 0; + gpt2_forward(&model, dummy_batch, NULL, B, T); + free(dummy_batch); + + // Setup logger and tokenizer + Logger logger; + logger_init(&logger, output_log_file); + + Tokenizer tokenizer; + tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); + + unsigned long long rng_state = 1337; + int *gen_tokens = (int *)mallocCheck(B * T * sizeof(int)); + float *cpu_logits = (float *)mallocCheck(model.config.vocab_size * sizeof(float)); + + // Training loop + struct timespec start, end; + double total_sum_iteration_time_s = 0.0; + + for (int step = 0; step <= train_num_batches; step++) { + int last_step = step == train_num_batches; + + // Validation + if (step % val_loss_every == 0 || last_step) { + float val_loss = 0.0f; + dataloader_reset(&val_loader); + for (int i = 0; i < val_num_batches; i++) { + dataloader_next_batch(&val_loader); + gpt2_forward(&model, val_loader.inputs, val_loader.targets, B, T); + val_loss += model.mean_loss; + } + val_loss /= val_num_batches; + if (my_pe == 1) { + printf("val loss %f\n", val_loss); + } + logger_log_val(&logger, step, val_loss); + } + + // Sampling (only on PE 1 which has the output) + if ((step > 0 && step % sample_every == 0) || last_step) { + for (int i = 0; i < B * T; ++i) gen_tokens[i] = GPT2_EOT; + if (my_pe == 1) printf("generating:\n---\n"); + for (int t = 1; t < genT; t++) { + gpt2_forward(&model, gen_tokens, NULL, B, T); + if (my_pe == 1) { + float *logits = model.acts.output + (t - 1) * model.config.padded_vocab_size; + cudaCheck(cudaMemcpy(cpu_logits, logits, model.config.vocab_size * sizeof(float), + cudaMemcpyDeviceToHost)); + float coin = random_f32(&rng_state); + int next_token = sample_softmax(cpu_logits, model.config.vocab_size, coin); + gen_tokens[t] = next_token; + if (tokenizer.init_ok) { + const char *token_str = tokenizer_decode(&tokenizer, next_token); + safe_printf(token_str); + } else { + printf("%d ", next_token); + } + fflush(stdout); + } + // Broadcast next token to all PEs + MPI_Bcast(&gen_tokens[t], 1, MPI_INT, 1, MPI_COMM_WORLD); + } + if (my_pe == 1) printf("\n---\n"); + } + + if (last_step) break; + + // Training step + clock_gettime(CLOCK_MONOTONIC, &start); + dataloader_next_batch(&train_loader); + gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T); + gpt2_zero_grad(&model); + gpt2_backward(&model); + gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step + 1); + cudaCheck(cudaDeviceSynchronize()); + 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; + + if (my_pe == 1) { + 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); + } + logger_log_train(&logger, step, model.mean_loss); + } + + if (my_pe == 0) { + printf("total average iteration time: %f ms\n", + total_sum_iteration_time_s / train_num_batches * 1000); + } + + // Cleanup + dataloader_free(&train_loader); + dataloader_free(&val_loader); + tokenizer_free(&tokenizer); + gpt2_free(&model); + free(cpu_logits); + free(gen_tokens); + cublasCheck(cublasDestroy(cublas_handle)); + logger_free(&logger); + + nvshmem_finalize(); + MPI_Finalize(); + + return 0; +} +#endif + diff --git a/test_nvshmem b/test_nvshmem deleted file mode 100755 index 9041d00..0000000 Binary files a/test_nvshmem and /dev/null differ