/* GPT-2 Transformer Neural Net trained in raw CUDA Non-trivial notes to be aware of: We are being clever in the backward pass to conserve memory. In particular, all parameters use a += in the backward pass, so we can later do gradient accumulation. But all activations have = instead of += because these are faster (just read, no write). This is okay for all activations except for those in the residual stream, where the gradients have to add. We make sure that those parts work out ok and that we do a += as necessary. E.g., the layernorms are connected to the residuals so we += in layernorm backward. In this file we are using Mixed Precision training, so different activations, paramaters, grads and buffers may be kept at different precisions, to take advantage of the fast low-precision hardware in the latest GPUs (bf16/fp16), and fp8 (coming soon^TM). Compile: make train_gpt2cu Example launch using bfloat16 on 1 GPU batch size 8, sample/eval every 200 steps: Also we're using TinyStories here for example as it is a bigger dataset ./train_gpt2cu -b 8 -v 200 -s 200 -i data/TinyStories Example launch using bfloat16 on 4 GPUs, same as above: mpirun -np 4 ./train_gpt2cu -b 8 -v 200 -s 200 -i data/TinyStories */ #include #include #include #include #include #include #include #include #include #include #include // GPU / CUDA related #include #include #include #include #include // Multi-GPU related #ifdef MULTI_GPU #include #include #endif // ---------------------------------------------------------------------------- // CUDA precision settings enum PrecisionMode { PRECISION_FP32, PRECISION_FP16, PRECISION_BF16 }; // fp32 #if defined(ENABLE_FP32) typedef float floatX; typedef float floatN; #define CUBLAS_LOWP CUDA_R_32F #define CUBLAS_LOWP_COMPUTE cublas_compute_type // auto-select FP32 vs TF32 const char* load_filename = "gpt2_124M.bin"; // fp32 weights PrecisionMode PRECISION_MODE = PRECISION_FP32; const char* precision_mode_str = "fp32"; #ifdef MULTI_GPU const ncclDataType_t ncclFloatX = ncclFloat; const ncclDataType_t ncclFloatN = ncclFloat; #endif // use fp16 (note: this may require gradient scaler, currently not implemented!) #elif defined(ENABLE_FP16) typedef half floatX; typedef float floatN; #define CUBLAS_LOWP CUDA_R_16F #define CUBLAS_LOWP_COMPUTE CUBLAS_COMPUTE_32F const char* load_filename = "gpt2_124M.bin"; // fp32 weights PrecisionMode PRECISION_MODE = PRECISION_FP16; const char* precision_mode_str = "fp16"; #ifdef MULTI_GPU const ncclDataType_t ncclFloatX = ncclHalf; const ncclDataType_t ncclFloatN = ncclFloat; #endif // bfloat16 (default!) #else typedef __nv_bfloat16 floatX; typedef float floatN; #define CUBLAS_LOWP CUDA_R_16BF #define CUBLAS_LOWP_COMPUTE CUBLAS_COMPUTE_32F const char* load_filename = "gpt2_124M_bf16.bin"; // bf16 weights PrecisionMode PRECISION_MODE = PRECISION_BF16; const char* precision_mode_str = "bf16"; #ifdef MULTI_GPU const ncclDataType_t ncclFloatX = ncclBfloat16; const ncclDataType_t ncclFloatN = ncclFloat; #endif #endif // ---------------------------------------------------------------------------- // CUDA utils // 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; static cublasComputeType_t cublas_compute_type; cublasHandle_t cublas_handle; cublasLtHandle_t cublaslt_handle; namespace cg = cooperative_groups; // convenience macro for calculating grid/block dimensions for kernels #define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) // CUDA error checking 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__)) // cuBLAS error checking 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__); } #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 #ifdef ENABLE_BF16 __device__ void atomicAddX(__nv_bfloat16* addr, __nv_bfloat16 val) { uintptr_t ptr_val = reinterpret_cast(addr); __nv_bfloat162* ptr_bf16 = reinterpret_cast<__nv_bfloat162*>(ptr_val & ~uintptr_t(0x3)); // Prepare the value to add, setting the other half to zero __nv_bfloat162 add_val = (ptr_val & 0x3) ? __halves2bfloat162(__ushort_as_bfloat16(0), val) : __halves2bfloat162(val, __ushort_as_bfloat16(0)); atomicAdd(ptr_bf16, add_val); } #endif #ifdef ENABLE_FP16 __device__ void atomicAddX(half* addr, half val) { uintptr_t ptr_val = reinterpret_cast(addr); half2* ptr_fp16 = reinterpret_cast(ptr_val & ~uintptr_t(0x3)); // Prepare the value to add, setting the other half to zero half2 add_val = (ptr_val & 0x3) ? __halves2half2(__ushort_as_half(0), val) : __halves2half2(val, __ushort_as_half(0)); atomicAdd(ptr_fp16, add_val); } #endif __device__ void atomicAddX(float* addr, float val) { atomicAdd(addr, val); } // ---------------------------------------------------------------------------- // Random Number Generatiom // Simple xorshift RNG __device__ __host__ unsigned int random_u32(unsigned long long *state) { // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A *state ^= *state >> 12; *state ^= *state << 25; *state ^= *state >> 27; return (*state * 0x2545F4914F6CDD1Dull) >> 32; } __device__ __host__ float random_f32(unsigned long long *state) { // random float32 in [0,1) return (random_u32(state) >> 8) / 16777216.0f; } // SquirrelNoise5 - Squirrel's Raw Noise utilities (version 5) // This gives us a random number from threadIdx/blockIdx + a single seed for the entire GPU // todo - possibly overkill and we don't need such high quality random numbers? (tbd) // http://eiserloh.net/noise/SquirrelNoise5.hpp __device__ __host__ constexpr unsigned int SquirrelNoise5(int positionX, unsigned int seed) { constexpr unsigned int SQ5_BIT_NOISE1 = 0xd2a80a3f; // 11010010101010000000101000111111 constexpr unsigned int SQ5_BIT_NOISE2 = 0xa884f197; // 10101000100001001111000110010111 constexpr unsigned int SQ5_BIT_NOISE3 = 0x6C736F4B; // 01101100011100110110111101001011 constexpr unsigned int SQ5_BIT_NOISE4 = 0xB79F3ABB; // 10110111100111110011101010111011 constexpr unsigned int SQ5_BIT_NOISE5 = 0x1b56c4f5; // 00011011010101101100010011110101 unsigned int mangledBits = (unsigned int) positionX; mangledBits *= SQ5_BIT_NOISE1; mangledBits += seed; mangledBits ^= (mangledBits >> 9); mangledBits += SQ5_BIT_NOISE2; mangledBits ^= (mangledBits >> 11); mangledBits *= SQ5_BIT_NOISE3; mangledBits ^= (mangledBits >> 13); mangledBits += SQ5_BIT_NOISE4; mangledBits ^= (mangledBits >> 15); mangledBits *= SQ5_BIT_NOISE5; mangledBits ^= (mangledBits >> 17); return mangledBits; } __device__ __host__ constexpr unsigned int Get1dNoiseUint(int positionX, unsigned int seed) { return SquirrelNoise5(positionX, seed); } __device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY, unsigned int seed) { constexpr int PRIME_NUMBER = 198491317; // Large prime number with non-boring bits return SquirrelNoise5(indexX + (PRIME_NUMBER * indexY), seed); } __device__ __host__ constexpr float Get1dNoiseZeroToOne(int index, unsigned int seed) { constexpr double ONE_OVER_MAX_UINT = (1.0 / (double) 0xFFFFFFFF); return (float)(ONE_OVER_MAX_UINT * (double) SquirrelNoise5(index, seed)); } __device__ __host__ constexpr float Get2dNoiseZeroToOne(int indexX, int indexY, unsigned int seed) { constexpr double ONE_OVER_MAX_UINT = (1.0 / (double) 0xFFFFFFFF); return (float)(ONE_OVER_MAX_UINT * (double) Get2dNoiseUint(indexX, indexY, seed)); } // stochastic rounding built on top of Squirel Noise above (with seed updated per step via xorshift) __device__ __forceinline__ void stochastic_rounding(float in, __nv_bfloat16 *out, unsigned int seed) { // todo - is this stochastic rounding *too good*? can we cut any corners? unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); unsigned int threshold = random & 0xFFFF; unsigned int float_bits = __float_as_uint(in); unsigned int rounded_bits = float_bits & 0x0000FFFF; float_bits = (rounded_bits > threshold) ? (float_bits | 0xFFFF) : (float_bits & ~0xFFFF); *out = __float2bfloat16_rn(__uint_as_float(float_bits)); } __device__ __forceinline__ void stochastic_rounding(float in, half *out, unsigned int random) { *out = (float)in; // todo - implement this... } __device__ __forceinline__ void stochastic_rounding(float in, float *out, unsigned int random) { *out = in; // dummy function for when floatX is float (FP32 mode) } // ---------------------------------------------------------------------------- // fread convenience utils, with nice handling of error checking using macros // simple replace fopen, fread, fclose with fopenCheck, freadCheck, fcloseCheck FILE *fopen_check(const char *path, const char *mode, const char *file, int line) { FILE *fp = fopen(path, mode); if (fp == NULL) { fprintf(stderr, "Error: Failed to open file '%s' at %s:%d\n", path, file, line); fprintf(stderr, "Error details:\n"); fprintf(stderr, " File: %s\n", file); fprintf(stderr, " Line: %d\n", line); fprintf(stderr, " Path: %s\n", path); fprintf(stderr, " Mode: %s\n", mode); fprintf(stderr, "---> HINT: try to re-run `python train_gpt2.py`\n"); exit(EXIT_FAILURE); } return fp; } #define fopenCheck(path, mode) fopen_check(path, mode, __FILE__, __LINE__) void fread_check(void *ptr, size_t size, size_t nmemb, FILE *stream, const char *file, int line) { size_t result = fread(ptr, size, nmemb, stream); if (result != nmemb) { if (feof(stream)) { fprintf(stderr, "Error: Unexpected end of file at %s:%d\n", file, line); } else if (ferror(stream)) { fprintf(stderr, "Error: File read error at %s:%d\n", file, line); } else { fprintf(stderr, "Error: Partial read at %s:%d. Expected %zu elements, read %zu\n", file, line, nmemb, result); } fprintf(stderr, "Error details:\n"); fprintf(stderr, " File: %s\n", file); fprintf(stderr, " Line: %d\n", line); fprintf(stderr, " Expected elements: %zu\n", nmemb); fprintf(stderr, " Read elements: %zu\n", result); exit(EXIT_FAILURE); } } #define freadCheck(ptr, size, nmemb, stream) fread_check(ptr, size, nmemb, stream, __FILE__, __LINE__) void fclose_check(FILE *fp, const char *file, int line) { if (fclose(fp) != 0) { fprintf(stderr, "Error: Failed to close file at %s:%d\n", file, line); fprintf(stderr, "Error details:\n"); fprintf(stderr, " File: %s\n", file); fprintf(stderr, " Line: %d\n", line); exit(EXIT_FAILURE); } } #define fcloseCheck(fp) fclose_check(fp, __FILE__, __LINE__) // ---------------------------------------------------------------------------- // malloc error-handling wrapper util void *malloc_check(size_t size, const char *file, int line) { void *ptr = malloc(size); if (ptr == NULL) { fprintf(stderr, "Error: Memory allocation failed at %s:%d\n", file, line); fprintf(stderr, "Error details:\n"); fprintf(stderr, " File: %s\n", file); fprintf(stderr, " Line: %d\n", line); fprintf(stderr, " Size: %zu bytes\n", size); exit(EXIT_FAILURE); } return ptr; } #define mallocCheck(size) malloc_check(size, __FILE__, __LINE__) // ---------------------------------------------------------------------------- // MPI / multi-processing setup // Parameters specific to training on multiple GPUs. typedef struct { int process_rank; // Rank of this process among all MPI processes. 0 if no multi-GPU. int num_processes; // Total number of processes. 1 if no multi-GPU. int local_device_idx; // This process GPU index on current machine. 0 if no multi-GPU. #ifdef MULTI_GPU ncclComm_t nccl_comm; // NCCL communication primitive, used for collective mutli-GPU work. #endif } MultiGpuConfig; // one global variable to hold the multi-GPU configuration for this process MultiGpuConfig multi_gpu_config; #ifdef MULTI_GPU // Determine which GPU this process should use. // Processes on the same machines use different GPU indicies. Processes on other machines don't. // Copied from NCCL examples: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/examples.html#example-2-one-device-per-process-or-thread int multi_gpu_get_local_device_idx(int process_rank, int num_processes) { char hostname[1024]; hostname[1023] = '\0'; // All processes on the same machine will share the same hostname. gethostname(hostname, 1023); for (int i=0; i < 1024; i++) { if (hostname[i] == '.') { hostname[i] = '\0'; break; } } uint64_t hostname_hash = 5381; for (int c = 0; hostname[c] != '\0'; c++){ hostname_hash = ((hostname_hash << 5) + hostname_hash) ^ hostname[c]; } // Distribute all hostname hashes to all processes. uint64_t* all_hostsname_hashes = (uint64_t*)malloc(num_processes * sizeof(uint64_t)); all_hostsname_hashes[process_rank] = hostname_hash; mpiCheck(MPI_Allgather(MPI_IN_PLACE, 0, MPI_DATATYPE_NULL, all_hostsname_hashes, sizeof(uint64_t), MPI_BYTE, MPI_COMM_WORLD)); // Identify which GPU we need to use. int local_device_idx = 0; for (int current_process = 0; current_process < num_processes; ++current_process) { if (current_process == process_rank) { // Found my gpu, local_device_idx now has my target GPU index. break; } if (all_hostsname_hashes[current_process] == all_hostsname_hashes[process_rank]) { // This process ID runs on the same machine, but it's not me, skip this GPU local_device_idx++; } } free(all_hostsname_hashes); return local_device_idx; } #endif MultiGpuConfig multi_gpu_config_init(int *argc, char ***argv) { #ifdef MULTI_GPU // Initialize MPI. MultiGpuConfig result; mpiCheck(MPI_Init(argc, argv)); mpiCheck(MPI_Comm_rank(MPI_COMM_WORLD, &result.process_rank)); mpiCheck(MPI_Comm_size(MPI_COMM_WORLD, &result.num_processes)); result.local_device_idx = multi_gpu_get_local_device_idx(result.process_rank, result.num_processes); cudaCheck(cudaSetDevice(result.local_device_idx)); ncclUniqueId nccl_id; if (result.process_rank == 0) { ncclCheck(ncclGetUniqueId(&nccl_id)); } mpiCheck(MPI_Bcast((void *)&nccl_id, sizeof(nccl_id), MPI_BYTE, 0, MPI_COMM_WORLD)); ncclCheck(ncclCommInitRank(&result.nccl_comm, result.num_processes, nccl_id, result.process_rank)); return result; #else printf("Multi-GPU support is disabled. Using a single GPU."); return MultiGpuConfig{ .process_rank = 0, .num_processes = 1, .local_device_idx = 0, }; #endif } void multi_gpu_config_free(const MultiGpuConfig* multi_gpu_config) { #ifdef MULTI_GPU ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm)); mpiCheck(MPI_Finalize()); #endif } // convenience function that only prints if the rank of process is zero void printf0(const char *format, ...) { if (multi_gpu_config.process_rank == 0) { va_list args; va_start(args, format); vprintf(format, args); va_end(args); } } // ---------------------------------------------------------------------------- // all the kernels // warp-level reduction for finding the maximum value __device__ float warpReduceMax(float val) { for (int offset = 16; offset > 0; offset /= 2) { val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset)); } return val; } // warp-level reduction for summing values __device__ float warpReduceSum(float val) { for (int offset = 16; offset > 0; offset /= 2) { val += __shfl_down_sync(0xFFFFFFFF, val, offset); } return val; } template __global__ void encoder_forward_kernel2(TOut* out, int* inp, Tw* wte, Tw* wpe, 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]; TOut* out_btc = out + b * T * C + t * C + c; Tw* wte_ix = wte + ix * C + c; Tw* wpe_tc = wpe + t * C + c; *out_btc = (TOut)((float)*wte_ix + (float)*wpe_tc); } } // really bad naive kernel with atomicAdd template __global__ void encoder_backward_kernel(Type* dwte, Type* dwpe, const Tdout* 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 Tdout* dout_btc = dout + b * T * C + t * C + c; Type* dwte_ix = dwte + ix * C + c; Type* dwpe_tc = dwpe + t * C + c; atomicAddX(dwte_ix, (Type)*dout_btc); atomicAddX(dwpe_tc, (Type)*dout_btc); } } // currently reads FP32, outputs floatX(FP16/BF16/FP8) template __global__ void layernorm_forward_kernel3(TOut* __restrict__ out, Type* __restrict__ mean, Type* __restrict__ rstd, const Type* __restrict__ inp, const TParam* __restrict__ weight, const TParam* __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; } // guard // the row of input that this group of threads is responsible for const Type* x = inp + idx * C; // mean float sum = 0.0f; for (int i = warp.thread_rank(); i < C; i += warp.size()) { sum += (float)x[i]; } sum = cg::reduce(warp, sum, cg::plus{}); float m = sum / C; if(warp.thread_rank() == 0 && mean != nullptr) { __stcs(mean + idx, (Type)m); } // rstd sum = 0.0f; for (int i = warp.thread_rank(); i < C; i += warp.size()) { float diff = (float)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, (Type)s); } // final normalization and scaling by weight/bias TOut* o = out + idx * C; for (int c = warp.thread_rank(); c < C; c += warp.size()) { // load and store using the .cs "streaming" hint to the compiler, // indicating that this data will not be reused soon, and can be streamed through the caches // this allows the threads to get more cache-hits for the (shared) weight and bias parameters float n = s * ((float)__ldcs(x+c) - m); __stcs(o+c, (TOut)(n * (float)weight[c] + (float)bias[c])); } } // inputs floatX, outputs FP32 (for current FP32-only activation path for this WIP) __global__ void permute_kernel(floatX* q, floatX* k, floatX* v, const floatX* inp, int B, int N, int NH, int d) { // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d) // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d) int idx = blockIdx.x * blockDim.x + threadIdx.x; // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_] 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(floatX* dinp, const floatX* dq, const floatX* dk, const floatX* 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(floatX* inp, floatX *out, int B, int N, int NH, int d) { // out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d) int idx = blockIdx.x * blockDim.x + threadIdx.x; // out[b][n][nh_][d_] <- inp[b][nh_][n][d_] 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(floatX* dinp, const floatX *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] = (floatX)dout[other_idx]; } } template __global__ void softmax_forward_kernel5(Type* out, float inv_temperature, const Type* inp, int N, int T) { // inp, out shape: (N, T, T), where N = B * NH // fuses the multiplication by scale inside attention // directly autoregressive, so we only compute the lower triangular part // uses the online softmax algorithm assert(T % 4 == 0); cg::thread_block block = cg::this_thread_block(); cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); // micro-optimization: we iterate backwards so that // after the softmax backward operation completes, the cache retains the // part of the matrix close to the upper left corner, which benefits the // matmul operation that immediately follows. // int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // forward order int idx = (gridDim.x - blockIdx.x - 1) * warp.meta_group_size() + warp.meta_group_rank(); // backward order if(idx >= N * T) { return; } int own_pos = idx % T; int pos_by_4 = own_pos / 4; // one row of inp, i.e. inp[idx, :] of shape (T,) const Type* x = inp + idx * T; // not INF, so we don't get NaNs accidentally when subtracting two values. float maxval = -FLT_MAX; float sumval = 0.0f; const Type* x_aligned = reinterpret_cast(__builtin_assume_aligned(x, 16)); for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) { float regarray[4]; #pragma unroll for (int k = 0; k < 4; ++k) { regarray[k] = (float)x_aligned[4*i + k]; } float old_maxval = maxval; for(int k = 0; k < 4; ++k) { maxval = fmaxf(maxval, regarray[k]); } sumval *= expf(inv_temperature * (old_maxval - maxval)); for(int k = 0; k < 4; ++k) { sumval += expf(inv_temperature * (regarray[k] - maxval)); } } if(4*pos_by_4 + warp.thread_rank() <= own_pos) { float old_maxval = maxval; maxval = fmaxf(maxval, (float)x[4*pos_by_4 + warp.thread_rank()]); sumval *= expf(inv_temperature * (old_maxval - maxval)); sumval += expf(inv_temperature * ((float)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; // divide the whole row by the sum for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { // recalculation is faster than doing the round-trip through memory. float ev = expf(inv_temperature * ((float)__ldcs(x + i) - global_maxval)); __stcs(out + idx * T + i, (Type)(ev * norm)); } } template __global__ void residual_forward_kernel(TOut* out, T1* inp1, T2* inp2, int N) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < N) { out[idx] = (TOut)((float)__ldcs(&inp1[idx]) + (float)__ldcs(&inp2[idx])); } } #define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI) __global__ void gelu_forward_kernel(floatX* out, const floatX* inp, int N) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < N) { float xi = (float)inp[i]; float cube = 0.044715f * xi * xi * xi; out[i] = (floatX)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)))); } } __global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floatX* dout, const int N) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < N) { float x = (float)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] = (floatX)(local_grad * (float)dout[i]); } } __global__ void softmax_forward_kernel7(float* out, const float* inp, int N, int C) { // out is (N, C) just like inp. Each row of inp will get softmaxed. // same as kernel4, but optimised for very large Cs with advanced unrolling // The trick is to read into a register array (all indices known at compile time) // and always read UNROLL_FACTOR values to maximise memory level parallelism // even if we would be out of bounds, we set the index to min(C-1, idx) // so we just do some unnecessary reads (obviously bad for small C) // the writes are in a separate loop with a conditional check for out of bounds // making it separate is necessary to convince the compiler to do the right thing const int UNROLL_FACTOR = 8; const int warpsPerBlock = blockDim.x / 32; extern __shared__ float shared[]; int idx = blockIdx.x; int tid = threadIdx.x; int warpId = threadIdx.x / 32; // warp index within a block int laneId = threadIdx.x % 32; // thread index within a warp // shared[] must be allocated to have 2 * warpsPerBlock elements // first half for max values, the second half for sum values float* maxvals = shared; float* sumvals = &shared[warpsPerBlock]; if (tid >= C) { maxvals[warpId] = -INFINITY; sumvals[warpId] = 0.0f; return; } const float* x = inp + idx * C; // input float* y = out + idx * C; // output // first, thread coarsening by directly accessing global memory in series float maxval = -INFINITY; for (int i = tid; i < C; i += blockDim.x * UNROLL_FACTOR) { #pragma unroll for (int u = 0; u < UNROLL_FACTOR; u++) { maxval = fmaxf(maxval, x[min(C - 1, i + u*blockDim.x)]); } } // now within-warp reductions for maxval maxval = warpReduceMax(maxval); // the 0th thread of each warp writes the maxval of that warp to shared memory if (laneId == 0) maxvals[warpId] = maxval; __syncthreads(); // now the 0th thread reduces the maxvals in shared memory, i.e. across warps if (tid == 0) { float val = maxvals[tid]; #pragma unroll for (int i = 1; i < warpsPerBlock; i++) { val = fmaxf(val, maxvals[i]); } // store the final max in the first position maxvals[0] = val; } __syncthreads(); // broadcast the max to all threads float offset = maxvals[0]; // compute expf and write the result to global memory // + thread coarsening for sum float sumval = 0.0f; for (int i = tid; i < C; i += blockDim.x * UNROLL_FACTOR) { float reg_array[UNROLL_FACTOR]; #pragma unroll for (int u = 0; u < UNROLL_FACTOR; u++) { reg_array[u] = __ldcs(&x[min(C - 1, i + u*blockDim.x)]); } #pragma unroll for (int u = 0; u < UNROLL_FACTOR; u++) { if (i + u*blockDim.x < C) { float output = expf(reg_array[u] - offset); y[min(C - 1, i + u*blockDim.x)] = output; // compiler likes redundant min()?! sumval += output; // combined into the same loop unlike kernel3 } } } // okay now we calculated exp(x - max(x)) // step 2: sum all the values and divide by the sum // within-warp reduction for sumval sumval = warpReduceSum(sumval); // write sumval to shared memory if (laneId == 0) sumvals[warpId] = sumval; __syncthreads(); // inter-thread reduction of sum if (tid == 0) { float val = sumvals[tid]; #pragma unroll for (int i = 1; i < warpsPerBlock; ++i) { val += sumvals[i]; } sumvals[0] = val; } __syncthreads(); // broadcast the sum to all threads float sum = sumvals[0]; // divide the whole row by the sum for (int i = tid; i < C; i += blockDim.x * UNROLL_FACTOR) { float reg_array[UNROLL_FACTOR]; #pragma unroll for (int u = 0; u < UNROLL_FACTOR; u++) { reg_array[u] = y[min(C - 1, i + u*blockDim.x)]; } #pragma unroll for (int u = 0; u < UNROLL_FACTOR; u++) { if (i + u*blockDim.x < C) { y[i + u*blockDim.x] = reg_array[u] / sum; } } } } // this kernel performs a column-wise reduction over dout, in PyTorch equivalent to: // dbias = dout.sum((0,1)) // the idea is to employ one block to reduce along several columns, // where each block has a width of 32 columns to ensure coalesced access. // at the end we accumulate the reductions performed by the warps in each block via shared memory template __global__ void matmul_backward_bias_kernel4(Td* dbias, const Td* dout, int B, int T, int OC) { // this kernel is launched with 1D grid_dim of OC/32 // for example let's say block_size is 128 extern __shared__ float smem[]; // of size block_size (128) const int warp_id = threadIdx.x / warpSize; // warp index in the block, 0,1,2,3 const int lane_id = threadIdx.x % warpSize; // thread index in the warp, 0,1,2,...,31 const int tl = blockIdx.x * warpSize; // pointer to the start column for this block const int vstep = blockDim.x / warpSize; // number of warps in a block, e.g. 4 // pointer to the start of the column for one lane of threads // so e.g. 4 threads (of the same lane_id) will reduce this one column const Td* dout_col = dout + tl + lane_id; // column reductions by looping through the rows // each of the 4 threads offsets by its warp_id and then skips by vstep // together these 4 threads cover all B*T rows of this (lane_id) column // importantly, consecutive threads (in threadId) are processing adjacent columns, // leading to a coalesced memory access pattern float dout_sum = 0.0f; for (int row = warp_id; row < B * T; row += vstep) { dout_sum += (float)dout_col[row * OC]; } smem[lane_id + warp_id * warpSize] = dout_sum; __syncthreads(); // warp_id 0 reduces the shared memory column-wise, linearly 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] = (Td)dout_sum; } } // uses shared memory instead for the reduces template __global__ void layernorm_backward_kernel2(Tdinp* dinp, Tparams* dweight, Tparams* dbias, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* rstd, int B, int T, int C) { extern __shared__ float shared[]; // size = 2 * C namespace cg = cooperative_groups; 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; } // thread guards int b = idx / T; int t = idx % T; const Tdout* dout_bt = dout + b * T * C + t * C; const Trest* inp_bt = inp + b * T * C + t * C; Tdinp* dinp_bt = dinp + b * T * C + t * C; const float mean_bt = (float)mean[b * T + t]; const float rstd_bt = (float)rstd[b * T + t]; // the first half of shared memory is bias, second is weight float* dbias_shared = shared; float* dweight_shared = shared + C; // init shared memory to zero #pragma unroll for(int i = threadIdx.x; i < C; i+= blockDim.x){ dbias_shared[i] = 0.0f; dweight_shared[i] = 0.0f; } __syncthreads(); // first: two reduce operations 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 = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)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; // now iterate again and accumulate all the gradients for (int i = warp.thread_rank(); i < C; i += warp.size()) { float norm_bti = ((float)inp_bt[i] - mean_bt) * rstd_bt; float dnorm_i = (float)weight[i] * (float)dout_bt[i]; // gradient contribution to bias atomicAdd(&dbias_shared[i], (float)dout_bt[i]); // gradient contribution to weight atomicAdd(&dweight_shared[i], norm_bti * (float)dout_bt[i]); // gradient contribution to input float dval = 0.0f; dval += dnorm_i; // term 1 dval -= dnorm_mean; // term 2 dval -= norm_bti * dnorm_norm_mean; // term 3 dval *= rstd_bt; // final scale dinp_bt[i] = (Tdinp)((float)dinp_bt[i] + dval); } __syncthreads(); // write to global memory for(int i = threadIdx.x; i < C; i+= blockDim.x) { atomicAddX(&dbias[i], (Tparams)dbias_shared[i]); atomicAddX(&dweight[i], (Tparams)dweight_shared[i]); } } __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const floatX* datt, const floatX* 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; // go through blocks in reverse order, so the slowest block starts first 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 floatX* att_bth = att + t * T; const floatX* datt_bth = datt + t * T; floatX* dpreatt_bth = dpreatt + t * T; float local_sum = 0; for (int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { local_sum += (float)att_bth[t2] * (float)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) { // don't touch the cache. Some parts will still be here from the previous loop, and // we want to exploit those. float acc = (float)__ldcs(att_bth + t3) * ((float)__ldcs(datt_bth + t3) - local_sum); __stcs(dpreatt_bth + t3, (floatX)(scale * acc)); } } } // Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation). // Reference: https://developer.nvidia.com/blog/lerp-faster-cuda __device__ inline float lerp(float start, float end, float weight) { return fma(weight, end, fma(-weight, start, start)); } // Termplate type T instead of floatx template __global__ void adamw_kernel3(Tp* params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters, float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay, unsigned int seed) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i >= num_parameters) return; // guard float grad = (float)grads_memory[i]; float m = m_memory[i]; float v = v_memory[i]; // update the first moment (momentum) m = lerp(grad, m, beta1); m_memory[i] = m; // update the second moment (RMSprop) v = lerp(grad * grad, v, beta2); v_memory[i] = v; m /= beta1_correction; // m_hat v /= beta2_correction; // v_hat // update the parameters (weight/bias) float param = (float)params_memory[i] - (learning_rate * (m / (sqrtf(v) + eps) + weight_decay * (float)params_memory[i])); unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x, seed); // todo - explain stochastic rounding here stochastic_rounding(param, ¶ms_memory[i], random); } struct SoftmaxParams { float Scale; float Offset; }; template __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_tile<32>& warp, int idx, const Type* inp, int V, int P) { // same but not float4 // one row of inp, i.e. inp[idx, :] of shape (V,) const Type* x = inp + idx * P; float thread_maxval = -INFINITY; float thread_sumval = 0.0f; // do the loop in reverse to maximise probability of L2 cache hits // so even small L2s get some hits on the 2nd read of the same thread for (int i = V + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) { float v = (float)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); } // two reductions of up to 1024 threads: // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) // this results in much cleaner assembly than a multi-warp cg::reduce __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; // reduce maxval within each warp float warp_maxval = cg::reduce(warp, thread_maxval, cg::greater{}); // thread 0 in each warp writes to shared memory if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; } __syncthreads(); // each thread now loads the maxval across previous warps // if the thread is "out of range" of data, use -FLT_MAX as the maxval warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX; // now reduce the maxval among the warp threads float block_maxval = cg::reduce(warp, warp_maxval, cg::greater{}); // each thread uses maxval to scale sumval to avoid numerical instability / overflow thread_sumval *= expf(thread_maxval - block_maxval); // (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory float warp_sumval = cg::reduce(warp, thread_sumval, cg::plus{}); if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; } __syncthreads(); // same strategy, now reduce sumval across warps warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f; float block_sumval = cg::reduce(warp, warp_sumval, cg::plus{}); // return the softmax parameters return SoftmaxParams{1.f / block_sumval, block_maxval}; } // same as 2 but not using float4 (see dev/cuda/classifier_fused.cu) // will _update_ logits to logit gradients template __global__ void fused_classifier_kernel3(Type* logits, Type* losses, Type* probs, const Type* dlosses, const int* targets, int B, int T, int V, int P) { namespace cg = cooperative_groups; 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]; // softmax (reading B * T * V, same logits read again below, hopefully still in cache) SoftmaxParams sp = prepare_softmax_blockwide_nofloat4(warp, idx, logits, V, P); // calculate the probability needed for the loss and update (single-threaded) if(threadIdx.x == 0) { float prob = expf((float)logits[idx * P + ix] - sp.Offset) * sp.Scale; losses[idx] = (Type)(-logf(prob)); } // very sensible default for dlosses is 1/(B*T), which is the uniform loss float dloss = dlosses != NULL ? (float)dlosses[idx] : 1.0f / (B*T); // calculate the gradients directly, saves bandwidth from probs during training // but also supports writing probs for inference-only and debugging const Type* logits_vec = logits + idx * P; for (int i = threadIdx.x; i < V; i += blockDim.x) { // this is the 2nd read of logits after the one in prepare_softmax2 // this data will never be needed again, so we reduce cache persistence float v = (float)__ldcs(&logits_vec[i]); float prob = expf(v - sp.Offset) * sp.Scale; if (probs != NULL) { probs[idx * P + i] = (Type)prob; } float indicator = (i == ix) ? 1.0f : 0.0f; logits[idx * P + i] = (Type)((prob - indicator) * dloss); } } // ---------------------------------------------------------------------------- // kernel launchers template void encoder_forward(TOut* out, int* inp, Tw* wte, Tw* wpe, 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_forward_kernel2<<>>(out, inp, wte, wpe, B, T, C); cudaCheck(cudaGetLastError()); } template void encoder_backward(Type* dwte, Type* dwpe, const Tdout* 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()); } template void layernorm_forward(TOut* out, Type* mean, Type* rstd, Type* inp, Tparam* weight, Tparam* 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()); } // uses cuBLAS void matmul_forward_cublas(floatX* out, floatX* inp, floatX* weight, floatX* bias, int B, int T, int C, int OC) { assert(bias == NULL); // bias is not supported for this kernel // FP16 alpha/beta need to be used if and only if CUBLAS_COMPUTE_16F const float alpha = 1.0f, beta = 0.0f; const half alpha_fp16 = (half)alpha, beta_fp16 = (half)beta; const void* alpha_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (const void*)&alpha_fp16 : (const void*)α const void* beta_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (const void*)&beta_fp16 : (const void*)β cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, OC, B*T, C, alpha_ptr, weight, CUBLAS_LOWP, C, inp, CUBLAS_LOWP, C, beta_ptr, out, CUBLAS_LOWP, OC, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); } // uses cuBLASLt to fuse the bias and gelu. does not work with OC = 50257 (last layer) // https://docs.nvidia.com/cuda/cublas/#cublasltmatmul // https://github.com/NVIDIA/CUDALibrarySamples/blob/master/cuBLASLt/LtSgemm/sample_cublasLt_LtSgemm.cu void matmul_forward_cublaslt(floatX* out, floatX* inp, floatX* weight, floatX* bias, int B, int T, int C, int OC) { int has_bias = (bias != NULL); // check bias alignment if(((uintptr_t)bias % 16) != 0) { printf("Bias pointer is not aligned (cuBLASLt requirement)!\n"); exit(EXIT_FAILURE); } // FP16 alpha/beta need to be used if and only if CUBLAS_COMPUTE_16F const float alpha = 1.0f, beta = 0.0f; const half alpha_fp16 = (half)alpha, beta_fp16 = (half)beta; const void* alpha_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (const void*)&alpha_fp16 : (const void*)α const void* beta_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (const void*)&beta_fp16 : (const void*)β int returnedResults = 0; cublasLtMatmulDesc_t operationDesc; cublasLtMatmulPreference_t preference; cublasLtMatrixLayout_t weightLayout; cublasLtMatrixLayout_t inputLayout; cublasLtMatrixLayout_t outputLayout; cublasLtMatrixLayout_t biasLayout; cublasLtMatmulHeuristicResult_t heuristic; // create the operation descriptor cublasOperation_t opNoTranspose = CUBLAS_OP_N; cublasOperation_t opTranspose = CUBLAS_OP_T; cublasLtEpilogue_t epilogueBias = CUBLASLT_EPILOGUE_BIAS; cudaDataType_t scale_type = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? CUDA_R_16F : CUDA_R_32F; cublasCheck(cublasLtMatmulDescCreate(&operationDesc, CUBLAS_LOWP_COMPUTE, scale_type)); cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opTranspose, sizeof(opTranspose))); cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opNoTranspose, sizeof(opNoTranspose))); cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogueBias, sizeof(epilogueBias))); cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias))); // define matrix layouts cublasCheck(cublasLtMatrixLayoutCreate(&weightLayout, CUBLAS_LOWP, C, OC, C)); cublasCheck(cublasLtMatrixLayoutCreate(&inputLayout, CUBLAS_LOWP, C, B*T, C)); cublasCheck(cublasLtMatrixLayoutCreate(&outputLayout, CUBLAS_LOWP, OC, B*T, OC)); cublasCheck(cublasLtMatrixLayoutCreate(&biasLayout, CUBLAS_LOWP, OC, 1, OC)); // create a preference handle with specified max workspace cublasCheck(cublasLtMatmulPreferenceCreate(&preference)); cublasCheck(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &cublaslt_workspace_size, sizeof(cublaslt_workspace_size))); // find a suitable algorithm cublasCheck(cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc, weightLayout, inputLayout, outputLayout, outputLayout, preference, 1, &heuristic, &returnedResults)); if (returnedResults == 0) { printf("No cuBLASLt algorithm: B: %d, T: %d, C: %d, OC: %d, bias: %d\n", B, T, C, OC, has_bias); exit(EXIT_FAILURE); } // call the matmul cublasCheck(cublasLtMatmul(cublaslt_handle, operationDesc, alpha_ptr, weight, weightLayout, inp, inputLayout, beta_ptr, out, outputLayout, out, outputLayout, &heuristic.algo, cublaslt_workspace, cublaslt_workspace_size, 0)); // cleanups cublasCheck(cublasLtMatmulPreferenceDestroy(preference)); cublasCheck(cublasLtMatmulDescDestroy(operationDesc)); cublasCheck(cublasLtMatrixLayoutDestroy(weightLayout)); cublasCheck(cublasLtMatrixLayoutDestroy(inputLayout)); cublasCheck(cublasLtMatrixLayoutDestroy(outputLayout)); cublasCheck(cublasLtMatrixLayoutDestroy(biasLayout)); } void attention_forward(floatX* out, floatX* qkvr, floatX* att, floatX* inp, int B, int T, int C, int NH) { // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer. // Its contents will be overwritten by this function. const int block_size = 256; const int softmax_block_size = 256; // inp is (B, T, 3C) QKV // preatt, att are (B, NH, T, T) // output is (B, T, C) int HS = C / NH; // head size // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS) floatX *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()); // IMPORTANT: alpha/beta are FP32 for CUBLAS_COMPUTE_32F even if FP16 inputs/outputs // But need FP16 scale for CUBLAS_COMPUTE_16F (no errors otherwise, just garbage results *sigh*) const float alpha = 1.0f; const float beta = 0.0f; const floatX alpha_lowp = (floatX)alpha; const floatX beta_lowp = (floatX)beta; void* alpha_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (void*)&alpha_lowp : (void*)α void* beta_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (void*)&beta_lowp : (void*)β floatX* preatt = inp; cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, alpha_ptr, k, CUBLAS_LOWP, HS, T * HS, q, CUBLAS_LOWP, HS, T * HS, beta_ptr, preatt, CUBLAS_LOWP, T, T * T, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); // multiply all elements of preatt elementwise by scale 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()); // new approach: first cuBLAS another batched matmul floatX* vaccum = inp; // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs) cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, alpha_ptr, v, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, beta_ptr, vaccum, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); // now unpermute // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side num_blocks = CEIL_DIV(B * T * C, block_size); unpermute_kernel<<>>(vaccum, out, B, T, NH, HS); cudaCheck(cudaGetLastError()); } template void residual_forward(TOut* out, T1* inp1, T2* 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(floatX* out, const floatX* 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(floatX* dinp, const floatX* inp, const floatX* 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 softmax_forward(float* out, float* inp, int N, int C) { int grid_size = N; const int block_size = 512; size_t shared_mem_size = 2 * block_size / 32 * sizeof(float); softmax_forward_kernel7<<>>(out, inp, N, C); cudaCheck(cudaGetLastError()); } void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias, floatX* dout, floatX* inp, floatX* weight, int B, int T, int C, int OC) { float one = 1.0f; float zero = 0.0f; // backward to input, uses = in the backward pass (set the gradient) cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one, weight, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &zero, dinp, CUBLAS_LOWP, C, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); // backward to weight, uses += in the backward pass (accumulate the gradient) cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one, inp, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &one, dweight, CUBLAS_LOWP, C, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP)); // backward to bias, if given, does a += if (dbias != NULL) { const int block_size = 1024; const int grid_size = OC / 32; // for now, OC must be divisible by 32 for this kernel to work matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); cudaCheck(cudaGetLastError()); } } template void layernorm_backward(Tdinp* dinp, Tparams* dweight, Tparams* dbias, const Tdout* dout, const Trest* inp, const Tparams* weight, const Trest* mean, const Trest* 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()); } // the sequence of transformations in this compound op is: // inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C) void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* datt, floatX* scratch, const floatX* dout, const floatX* qkvr, const floatX* att, int B, int T, int C, int NH) { const int block_size = 256; int HS = C / NH; // head size // FP16 alpha/beta need to be used if and only if CUBLAS_COMPUTE_16F const float alpha = 1.0f, beta = 0.0f; const half alpha_fp16 = (half)alpha, beta_fp16 = (half)beta; const void* alpha_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (const void*)&alpha_fp16 : (const void*)α const void* beta_ptr = (CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F) ? (const void*)&beta_fp16 : (const void*)β // unpack convenience pointers into q, k, v const floatX *q, *k, *v; q = qkvr + 0 * B * T * C; k = qkvr + 1 * B * T * C; v = qkvr + 2 * B * T * C; floatX *dq, *dk, *dv; dq = dqkvr + 0 * B * T * C; dk = dqkvr + 1 * B * T * C; dv = dqkvr + 2 * B * T * C; // backward through the unpermute operation int num_blocks = CEIL_DIV(B * T * C, block_size); unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); cudaCheck(cudaGetLastError()); // backward into datt cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, alpha_ptr, v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, beta_ptr, datt, CUBLAS_LOWP, T, T * T, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); // backward into dv cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, alpha_ptr, scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, beta_ptr, dv, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); // backward into preatt int hs = C / NH; // head size float scale = 1.0f / sqrtf(hs); softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); cudaCheck(cudaGetLastError()); // backward into q cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, alpha_ptr, k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, beta_ptr, dq, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); // backward into k cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, alpha_ptr, q, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, beta_ptr, dk, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT)); // backward into inp num_blocks = CEIL_DIV(B * NH * T * HS, block_size); permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); cudaCheck(cudaGetLastError()); } // replaces logits with logit gradients template void fused_classifier3(Type* logits, Type* losses, const Type* 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, (Type*)NULL, dlosses, targets, B, T, V, P); cudaCheck(cudaGetLastError()); } // ---------------------------------------------------------------------------- // GPT-2 model definition typedef struct { int max_seq_len; // max sequence length, e.g. 1024 int vocab_size; // vocab size, e.g. 50257 int num_layers; // number of layers, e.g. 12 int num_heads; // number of heads in attention, e.g. 12 int channels; // number of channels, e.g. 768 } GPT2Config; // the parameters of the model // note the layernorms are kept in higher precision (floatN) constexpr const int NUM_PARAMETER_TENSORS = 16; typedef struct { floatX* wte; // (V, C) floatX* wpe; // (maxT, C) floatN* ln1w; // (L, C) floatN* ln1b; // (L, C) floatX* qkvw; // (L, 3*C, C) floatX* qkvb; // (L, 3*C) floatX* attprojw; // (L, C, C) floatX* attprojb; // (L, C) floatN* ln2w; // (L, C) floatN* ln2b; // (L, C) floatX* fcw; // (L, 4*C, C) floatX* fcb; // (L, 4*C) floatX* fcprojw; // (L, C, 4*C) floatX* fcprojb; // (L, C) floatN* lnfw; // (C) floatN* lnfb; // (C) } ParameterTensors; static_assert(sizeof(ParameterTensors) == NUM_PARAMETER_TENSORS * sizeof(void*), "Inconsistent sizes!"); void fill_in_parameter_sizes(size_t* param_sizes, size_t* param_sizeof, GPT2Config config) { size_t V = config.vocab_size; size_t C = config.channels; size_t maxT = config.max_seq_len; size_t L = config.num_layers; param_sizes[0] = V * C; // wte param_sizes[1] = maxT * C; // wpe param_sizes[2] = L * C; // ln1w param_sizes[3] = L * C; // ln1b param_sizes[4] = L * (3 * C) * C; // qkvw param_sizes[5] = L * (3 * C); // qkvb param_sizes[6] = L * C * C; // attprojw param_sizes[7] = L * C; // attprojb param_sizes[8] = L * C; // ln2w param_sizes[9] = L * C; // ln2b param_sizes[10] = L * (4 * C) * C; // fcw param_sizes[11] = L * (4 * C); // fcb param_sizes[12] = L * C * (4 * C); // fcprojw param_sizes[13] = L * C; // fcprojb param_sizes[14] = C; // lnfw param_sizes[15] = C; // lnfb // populate the parameter sizes in bytes for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { param_sizeof[i] = sizeof(floatX); } // override layernorms here below param_sizeof[2] = sizeof(floatN); // ln1w param_sizeof[3] = sizeof(floatN); // ln1b param_sizeof[8] = sizeof(floatN); // ln2w param_sizeof[9] = sizeof(floatN); // ln2b param_sizeof[14] = sizeof(floatN); // lnfw param_sizeof[15] = sizeof(floatN); // lnfb } // allocate memory for the parameters and point the individual tensors to the right places void* malloc_and_point_parameters(ParameterTensors* params, size_t* param_elements, size_t *param_sizeof) { // calculate the total number of parameters and bytes across all tensors size_t num_parameters = 0; size_t num_parameters_bytes = 0; for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { num_parameters += param_elements[i]; num_parameters_bytes += param_elements[i] * param_sizeof[i]; } // malloc all parameters all at once on the device void* params_memory; cudaCheck(cudaMalloc((void**)¶ms_memory, num_parameters_bytes)); // assign all the tensors their place in the array floatX** ptrs[] = { ¶ms->wte, ¶ms->wpe, (floatX**)¶ms->ln1w, (floatX**)¶ms->ln1b, ¶ms->qkvw, ¶ms->qkvb, ¶ms->attprojw, ¶ms->attprojb, (floatX**)¶ms->ln2w, (floatX**)¶ms->ln2b, ¶ms->fcw, ¶ms->fcb, ¶ms->fcprojw, ¶ms->fcprojb, (floatX**)¶ms->lnfw, (floatX**)¶ms->lnfb }; char* params_memory_iterator = (char*)params_memory; for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { *(ptrs[i]) = (floatX*)params_memory_iterator; params_memory_iterator += param_elements[i] * param_sizeof[i]; } return params_memory; } #define NUM_ACTIVATION_TENSORS 21 typedef struct { floatX* encoded; // (B, T, C) floatX* ln1; // (L, B, T, C) floatX* ln1_mean; // (L, B, T) floatX* ln1_rstd; // (L, B, T) floatX* atty; // (L, B, T, C) floatX* att; // (L, B, NH, T, T) floatX* attproj; // (L, B, T, C) floatX* residual2; // (L, B, T, C) floatX* ln2; // (L, B, T, C) floatX* ln2_mean; // (L, B, T) floatX* ln2_rstd; // (L, B, T) floatX* fch; // (L, B, T, 4*C) floatX* fch_gelu; // (L, B, T, 4*C) floatX* fcproj; // (L, B, T, C) floatX* residual3; // (L, B, T, C) floatX* lnf; // (B, T, C) floatX* lnf_mean; // (B, T) floatX* lnf_rstd; // (B, T) floatX* losses; // (B, T) // adding these two compared to the CPU .c code, needed for attention kernel as buffers floatX* qkvr; // (L, B, T, 3*C) // in inference mode, this buffer will store the logits // in training mode, this buffer will contain the *gradients* of the logits. // during the processing of transformer blocks, we will also use this as a // general scratchpad buffer. Allocation is made large enough to hold (B, T, 3C), // (B, NH, T, T), and (B, T, V) shaped tensors. floatX* output; } ActivationTensors; void fill_in_activation_sizes(size_t* act_sizes, size_t B, size_t T, GPT2Config config) { size_t V = config.vocab_size; size_t L = config.num_layers; size_t NH = config.num_heads; size_t C = config.channels; act_sizes[0] = B * T * C; // encoded act_sizes[1] = L * B * T * C; // ln1 act_sizes[2] = L * B * T; // ln1_mean act_sizes[3] = L * B * T; // ln1_rstd act_sizes[4] = L * B * T * C; // atty act_sizes[5] = L * B * NH * T * T; // att act_sizes[6] = L * B * T * C; // attproj act_sizes[7] = L * B * T * C; // residual2 act_sizes[8] = L * B * T * C; // ln2 act_sizes[9] = L * B * T; // ln2_mean act_sizes[10] = L * B * T; // ln2_rstd act_sizes[11] = L * B * T * 4*C; // fch act_sizes[12] = L * B * T * 4*C; // fch_gelu act_sizes[13] = L * B * T * C; // fcproj act_sizes[14] = L * B * T * C; // residual3 act_sizes[15] = B * T * C; // lnf act_sizes[16] = B * T; // lnf_mean act_sizes[17] = B * T; // lnf_rstd act_sizes[18] = B * T; // losses act_sizes[19] = L * B * T * 3*C; // qkvr act_sizes[20] = B * T * max(3*C, max(NH*T, V)); // output / scratch } // Backward pass is conceptually quite different from forward, because we can discard // the activations of a layer as soon as we're done with it. This lets us aggressively // reuse memory, so that we need far fewer tensors for backward state. #define NUM_BACKWARD_TENSORS 3 typedef struct { floatX* bt4c; // (B, T, 4*C) floatX* preatt; // (B, NH, T, T) floatX* residual3; // (B, T, C) } GradActTensors; void fill_in_grad_act_sizes(size_t* act_sizes, size_t B, size_t T, GPT2Config config) { size_t NH = config.num_heads; size_t C = config.channels; act_sizes[0] = B * T * 4 * C; // bt4c act_sizes[1] = B * NH * T * T; // preatt act_sizes[2] = B * T * C; // residual3 } void* malloc_and_point(floatX** targets[], const size_t* act_sizes, size_t n) { size_t num_activations = 0; for (size_t i = 0; i < n; i++) { num_activations += act_sizes[i]; } void* acts_memory; cudaCheck(cudaMalloc((void**)&acts_memory, num_activations * sizeof(floatX))); char* acts_memory_iterator = (char*)acts_memory; for (size_t i = 0; i < n; i++) { *(targets[i]) = (floatX*)acts_memory_iterator; acts_memory_iterator += act_sizes[i] * sizeof(floatX); } return acts_memory; } void* malloc_and_point_activations(ActivationTensors* acts, const size_t* act_sizes) { floatX** 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 }; return malloc_and_point(ptrs, act_sizes, NUM_ACTIVATION_TENSORS); } void* malloc_and_point_backward(GradActTensors* acts, const size_t* act_sizes) { floatX** ptrs[] = { &acts->bt4c, &acts->preatt, &acts->residual3 }; return malloc_and_point(ptrs, act_sizes, NUM_BACKWARD_TENSORS); } typedef struct { GPT2Config config; // the weights of the model, and their sizes ParameterTensors params; size_t param_elements[NUM_PARAMETER_TENSORS]; size_t param_sizeof[NUM_PARAMETER_TENSORS]; void* params_memory; size_t num_parameters; size_t num_parameters_bytes; // gradients of the weights ParameterTensors grads; void* grads_memory; // buffers for the AdamW optimizer float* m_memory; float* v_memory; // the activations of the model, and their sizes ActivationTensors acts; size_t act_sizes[NUM_ACTIVATION_TENSORS]; void* acts_memory; size_t num_activations; // gradients of the activations GradActTensors grads_acts; size_t num_grad_acts; void* grads_acts_memory; // other run state configuration int batch_size; // the batch size (B) of current forward pass int seq_len; // the sequence length (T) of current forward pass 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; void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { if (PRECISION_MODE == PRECISION_FP16) { // TODO for later perhaps, would require us dynamically converting the // model weights from fp32 to fp16 online, here in this function, or writing // the fp16 weights directly from Python, which we only do for fp32/bf16 atm. fprintf(stderr, "build_from_checkpoint() does not support fp16 right now.\n"); exit(EXIT_FAILURE); } // read in model from a checkpoint file FILE *model_file = fopenCheck(checkpoint_path, "rb"); int model_header[256]; freadCheck(model_header, sizeof(int), 256, model_file); if (model_header[0] != 20240326) { printf("Bad magic model file\n"); exit(EXIT_FAILURE); } int version = model_header[1]; if (!(version == 1 || version == 2)) { // 1 = fp32, ordered layernorm at the end // 2 = bf16, ordered layernorm at the end fprintf(stderr, "Bad version in model file\n"); fprintf(stderr, "---> HINT: try to re-run `python train_gpt2.py`\n"); exit(EXIT_FAILURE); } // read in hyperparameters model->config.max_seq_len = model_header[2]; model->config.vocab_size = model_header[3]; model->config.num_layers = model_header[4]; model->config.num_heads = model_header[5]; model->config.channels = model_header[6]; // allocate space for all the parameters and read them in fill_in_parameter_sizes(model->param_elements, model->param_sizeof, model->config); model->num_parameters = 0; model->num_parameters_bytes = 0; for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { model->num_parameters += model->param_elements[i]; model->num_parameters_bytes += model->param_elements[i] * model->param_sizeof[i]; } // create memory for model parameters on the device model->params_memory = malloc_and_point_parameters(&model->params, model->param_elements, model->param_sizeof); // read in all the parameters from file and copy them to device float* params_memory_cpu = (float*)mallocCheck(model->num_parameters_bytes); freadCheck(params_memory_cpu, 1, model->num_parameters_bytes, model_file); cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, model->num_parameters_bytes, cudaMemcpyHostToDevice)); free(params_memory_cpu); fcloseCheck(model_file); // other inits 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; // -1.0f will designate no loss model->rng_state = 13371337; } void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { // targets are optional and could be NULL // in this function we must be careful and use size_t instead of int, otherwise // we could overflow int. E.g. l * B * NH * T * T overflows int at B 16. // ensure the model was initialized or error out if (model->params_memory == NULL) { printf("Error: model was not initialized properly.\n"); exit(EXIT_FAILURE); } // convenience parameters size_t V = model->config.vocab_size; size_t L = model->config.num_layers; size_t NH = model->config.num_heads; size_t C = model->config.channels; // validate inputs, all indices must be in the range [0, V) 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); } } // allocate space for all the activations if needed (done here, lazily) if(model->acts_memory == NULL) { // record the current B,T as well model->batch_size = B; model->seq_len = T; // allocate the space fill_in_activation_sizes(model->act_sizes, B, T, model->config); 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); printf0("allocated %d MiB for activations\n", (int)round(num_activations * sizeof(floatX) / (1024 * 1024))); // also create memory for caching inputs and targets cudaCheck(cudaMalloc((void**)&model->inputs, B * T * sizeof(int))); cudaCheck(cudaMalloc((void**)&model->targets, B * T * sizeof(int))); cudaCheck(cudaMallocHost((void**)&model->cpu_losses, B * T * sizeof(floatX))); } else { // validate B,T is consistent with how we've allocated the memory before // in principle we could get more clever here in the future, for now this is safest 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, (int)B, (int)T); exit(EXIT_FAILURE); } } // copy inputs/targets to the model cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice)); if (targets != NULL) { cudaCheck(cudaMemcpy(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice)); } // forward pass ParameterTensors params = model->params; // for brevity ActivationTensors acts = model->acts; floatX* residual; encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0] for (int l = 0; l < L; l++) { residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; // get the pointers of the weights for this layer floatN* l_ln1w = params.ln1w + l * C; floatN* l_ln1b = params.ln1b + l * C; floatX* l_qkvw = params.qkvw + l * 3*C * C; floatX* l_qkvb = params.qkvb + l * 3*C; floatX* l_attprojw = params.attprojw + l * C * C; floatX* l_attprojb = params.attprojb + l * C; floatN* l_ln2w = params.ln2w + l * C; floatN* l_ln2b = params.ln2b + l * C; floatX* l_fcw = params.fcw + l * 4*C * C; floatX* l_fcb = params.fcb + l * 4*C; floatX* l_fcprojw = params.fcprojw + l * C * 4*C; floatX* l_fcprojb = params.fcprojb + l * C; // get the pointers of the activations for this layer floatX* l_ln1 = acts.ln1 + l * B * T * C; floatX* l_ln1_mean = acts.ln1_mean + l * B * T; floatX* l_ln1_rstd = acts.ln1_rstd + l * B * T; floatX* l_qkvr = acts.qkvr + l * B * T * 3*C; floatX* l_atty = acts.atty + l * B * T * C; floatX* l_att = acts.att + l * B * NH * T * T; floatX* l_attproj = acts.attproj + l * B * T * C; floatX* l_residual2 = acts.residual2 + l * B * T * C; floatX* l_ln2 = acts.ln2 + l * B * T * C; floatX* l_ln2_mean = acts.ln2_mean + l * B * T; floatX* l_ln2_rstd = acts.ln2_rstd + l * B * T; floatX* l_fch = acts.fch + l * B * T * 4*C; floatX* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C; floatX* l_fcproj = acts.fcproj + l * B * T * C; floatX* l_residual3 = acts.residual3 + l * B * T * C; // these are only needed as scratchpads for the forward pass, but // need not be stored for backward floatX* scratch = (floatX*)acts.output; // now do the forward pass layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C); matmul_forward_cublaslt(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_cublaslt(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_cublaslt(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_cublaslt(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 + (L-1) * B * T * C; // last residual is in residual3 layernorm_forward(acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, params.lnfw, params.lnfb, B, T, C); matmul_forward_cublas(acts.output, acts.lnf, params.wte, NULL, B, T, C, V); // also forward the cross-entropy loss function if we have the targets if (targets != NULL) { // fused classifier: does the forward pass and first part of the backward pass // we're passing dlosses = NULL, which will default them to 1.0f/(B*T), i.e. uniform loss fused_classifier3(acts.output, acts.losses, (floatX*)NULL, model->targets, B, T, V, V); // for convenience also evaluate the mean loss (TODO re-think this compute+sync point) // move the (B,T) losses to CPU cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(floatX), cudaMemcpyDeviceToHost)); float mean_loss = 0.0f; for (int i=0; icpu_losses[i]); } mean_loss /= B*T; model->mean_loss = mean_loss; } else { // if we don't have targets, we don't have loss model->mean_loss = -1.0f; } } void gpt2_zero_grad(GPT2 *model) { if (model->grads_acts_memory != NULL) { cudaCheck(cudaMemset(model->grads_acts_memory, 0, model->num_grad_acts * sizeof(floatX))); } if (model->grads_memory != NULL) { cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX))); } } void gpt2_backward(GPT2 *model) { // double check we forwarded previously, with targets if (model->mean_loss == -1.0f) { printf("Error: must forward with targets before backward\n"); exit(EXIT_FAILURE); } // lazily allocate the memory for gradients of the weights and activations, if needed if (model->grads_memory == NULL) { // allocate buffers for weight gradients model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_elements, model->param_sizeof); printf0("allocated %d MiB for parameter gradients\n", (int)round(model->num_parameters * sizeof(floatX) / (1024 * 1024))); // we're going to be clever for the activations backward pass. we don't need to exactly // mirror the forward pass activations and we will save memory. size_t bw_act_sizes[NUM_ACTIVATION_TENSORS]; fill_in_grad_act_sizes(bw_act_sizes, model->batch_size, model->seq_len, model->config); // count up and allocate the space model->grads_acts_memory = malloc_and_point_backward(&model->grads_acts, bw_act_sizes); model->num_grad_acts = 0; for (size_t i = 0; i < NUM_BACKWARD_TENSORS; i++) { model->num_grad_acts += bw_act_sizes[i]; } printf0("allocated %d MiB for activation gradients\n", (int)round(model->num_grad_acts * sizeof(floatX) / (1024 * 1024))); // init gradients of parameters and activations to zero gpt2_zero_grad(model); } // convenience shortcuts, size_t instead of int so that pointer arithmetics don't overflow size_t B = model->batch_size; size_t T = model->seq_len; size_t V = model->config.vocab_size; size_t L = model->config.num_layers; size_t NH = model->config.num_heads; size_t C = model->config.channels; // backward pass: go in the reverse order of the forward pass, and call backward() functions ParameterTensors params = model->params; // for brevity ParameterTensors grads = model->grads; ActivationTensors acts = model->acts; GradActTensors grads_acts = model->grads_acts; // we kick off the chain rule by filling in dlosses with 1.0f/(B*T) // this was done in the fused classifier kernel as last step of forward pass // technically that is a small, inline backward() pass of calculating // total, final loss as the mean over all losses over all (B,T) positions in the batch // next: backward the classifier matmul matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, B, T, C, V); // backward the final layernorm floatX* residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 floatX* dresidual = (floatX*)grads_acts.residual3; // the main buffer holding the gradient in the backward pass layernorm_backward(dresidual, grads.lnfw, grads.lnfb, grads_acts.bt4c, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C); // now backward all the layers for (int l = L-1; l >= 0; l--) { residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; // get the pointers of the weights for this layer floatN* l_ln1w = params.ln1w + l * C; floatX* l_qkvw = params.qkvw + l * 3*C * C; floatX* l_attprojw = params.attprojw + l * C * C; floatN* l_ln2w = params.ln2w + l * C; floatX* l_fcw = params.fcw + l * 4*C * C; floatX* l_fcprojw = params.fcprojw + l * C * 4*C; // get the pointers of the gradients of the weights for this layer floatN* dl_ln1w = grads.ln1w + l * C; floatN* dl_ln1b = grads.ln1b + l * C; floatX* dl_qkvw = grads.qkvw + l * 3*C * C; floatX* dl_qkvb = grads.qkvb + l * 3*C; floatX* dl_attprojw = grads.attprojw + l * C * C; floatX* dl_attprojb = grads.attprojb + l * C; floatN* dl_ln2w = grads.ln2w + l * C; floatN* dl_ln2b = grads.ln2b + l * C; floatX* dl_fcw = grads.fcw + l * 4*C * C; floatX* dl_fcb = grads.fcb + l * 4*C; floatX* dl_fcprojw = grads.fcprojw + l * C * 4*C; floatX* dl_fcprojb = grads.fcprojb + l * C; // get the pointers of the activations for this layer floatX* l_ln1 = acts.ln1 + l * B * T * C; floatX* l_ln1_mean = acts.ln1_mean + l * B * T; floatX* l_ln1_rstd = acts.ln1_rstd + l * B * T; floatX* l_qkvr = acts.qkvr + l * B * T * 3*C; floatX* l_atty = acts.atty + l * B * T * C; floatX* l_att = acts.att + l * B * NH * T * T; floatX* l_residual2 = acts.residual2 + l * B * T * C; floatX* l_ln2 = acts.ln2 + l * B * T * C; floatX* l_ln2_mean = acts.ln2_mean + l * B * T; floatX* l_ln2_rstd = acts.ln2_rstd + l * B * T; floatX* l_fch = acts.fch + l * B * T * 4*C; floatX* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C; // get the pointers of the gradients of the activations for this layer // notice that there is no l *, because we just have a single copy, and keep // re-using this memory in every Transformer block as we calculate backward pass // we need a B x T x C buffer; thankfully, the forward activation for lnf isn't needed anymore, // so we can co-opt it here. floatX* dl_btc = (floatX*)acts.lnf; floatX* dl_bt4c = (floatX*)grads_acts.bt4c; floatX* dl_preatt = (floatX*)grads_acts.preatt; // re-use scratch buffer of the forward pass floatX* scratch = (floatX*)acts.output; // backprop this layer 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 does += to the dresidual, so it correctly accumulates grad from the MLP block above 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); // we more B x T x (4)C buffers. l_atty and l_fch aren't needed anymore at this point, so reuse their memory floatX* buffer_a = l_atty; floatX* buffer_b = l_fch; // this is B x T x 4C, so even larger than what we need 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 does += to dresidual, so it correctly accumulates gradient for the Attention block above 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); } // 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 // lazily allocate the memory for m_memory and v_memory 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))); printf0("allocated %d MiB for AdamW optimizer state m\n", (int)round(model->num_parameters * sizeof(float) / (1024 * 1024))); printf0("allocated %d MiB for AdamW optimizer state v\n", (int)round(model->num_parameters * sizeof(float) / (1024 * 1024))); } int block_size = 512; float beta1_correction = 1.0f - powf(beta1, t); float beta2_correction = 1.0f - powf(beta2, t); // Adam upadte // We need to know the parameter types (float or floatX) to process consecutive chunks char* params_mem = (char*)model->params_memory; char* grads_mem = (char*)model->grads_memory; size_t num_elements = model->param_elements[0]; size_t last_sizeof = model->param_sizeof[0]; size_t current_element = 0; for (int i = 1; i <= NUM_PARAMETER_TENSORS; i++) { if (i == NUM_PARAMETER_TENSORS || model->param_sizeof[i] != last_sizeof) { unsigned int seed = random_u32(&model->rng_state); // seed for stochastic rounding int num_blocks = CEIL_DIV(num_elements, block_size); // atm some params are in low precision (floatX) and some are in high precision (float) if (last_sizeof == sizeof(floatX)) { adamw_kernel3<<>>((floatX*)params_mem, (floatX*)grads_mem, &model->m_memory[current_element], &model->v_memory[current_element], num_elements, learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, seed); } else { adamw_kernel3<<>>((float*)params_mem, (float*)grads_mem, &model->m_memory[current_element], &model->v_memory[current_element], num_elements, learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, seed); } params_mem += num_elements * last_sizeof; grads_mem += num_elements * last_sizeof; current_element += num_elements; num_elements = 0; } if (i != NUM_PARAMETER_TENSORS) { num_elements += model->param_elements[i]; last_sizeof = model->param_sizeof[i]; } } 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); } #ifndef TESTING // if we are TESTING (see test_gpt2.cu), we'll skip the int main below // ---------------------------------------------------------------------------- // 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. use size_t to prevent overflow size_t B; size_t T; // input handling and its state FILE* tokens_file; long file_size; long current_position; // output memory int* batch; int* inputs; int* targets; // convenience variables size_t num_batches; } DataLoader; void dataloader_init(DataLoader *loader, const MultiGpuConfig* multi_gpu_config, const char* filename, size_t B, size_t T) { loader->process_rank = multi_gpu_config->process_rank; loader->num_processes = multi_gpu_config->num_processes; loader->B = B; loader->T = T; // open the input file for reading loader->tokens_file = fopenCheck(filename, "rb"); // determine the file size fseek(loader->tokens_file, 0, SEEK_END); loader->file_size = ftell(loader->tokens_file); fseek(loader->tokens_file, 0, SEEK_SET); if (loader->file_size < (B * T + 1) * sizeof(int)) { printf("Error: file size is too small for the batch size and sequence length\n"); exit(EXIT_FAILURE); } 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 // See: https://developer.nvidia.com/blog/how-optimize-data-transfers-cuda-cc/ 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 / (loader->num_processes * B * T * sizeof(int)); } void dataloader_reset(DataLoader *loader) { loader->current_position = 0; } void dataloader_next_batch(DataLoader *loader) { size_t B = loader->B; size_t T = loader->T; // if we are at the end of the file, loop back to the beginning 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*num_processes integers loader->current_position += loader->num_processes * B * T * sizeof(int); } void dataloader_free(DataLoader *loader) { fcloseCheck(loader->tokens_file); cudaFreeHost(loader->batch); } // ---------------------------------------------------------------------------- // sampler: takes probabilities and samples integers from them #define GPT2_EOT 50256 int sample_softmax(const float* logits, int n, float coin) { // sample index from logits (converted to probabilities using softmax) // coin is a random number in [0, 1), usually from random_f32() double norm = 0; for (int i = 0; i < n; i++) { norm += expf(logits[i]); } // instead of dividing all exp(logits), we can just multiply coin. 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; // in case of rounding errors } // ---------------------------------------------------------------------------- // Tokenizer (only supports decoding: tokens (integers) -> strings) typedef struct { uint32_t vocab_size; char **token_table; int init_ok; } Tokenizer; void safe_printf(const char *piece) { // the tokens are raw bytes, and we we only want to print the printable ones // many bytes can be various control codes, backspace, etc. if (piece == NULL) { return; } if (piece[0] == '\0') { return; } // handle individual byte tokens // every token is asserted to be at least one byte so doing piece[1] is ok if (piece[1] == '\0') { unsigned char byte_val = piece[0]; if (!(isprint(byte_val) || isspace(byte_val))) { return; // weird byte, don't print it } } printf("%s", piece); } void tokenizer_init(Tokenizer *tokenizer, const char *filename) { FILE *file = fopen(filename, "rb"); if (file == NULL) { // try to be more helpful as we just added this feature, erase later printf("---\n"); printf("WARNING: Failed to open the tokenizer file %s\n", filename); printf("The Tokenizer is a new feature added April 14 2024.\n"); printf("Re-run `python train_gpt2.py` to write it\n"); printf("---\n"); tokenizer->init_ok = 0; return; } // read in the header uint32_t header[256]; freadCheck(header, sizeof(uint32_t), 256, file); assert(header[0] == 20240328); assert(header[1] == 1); tokenizer->vocab_size = header[2]; // read in all the tokens unsigned char length; tokenizer->token_table = (char **)mallocCheck(tokenizer->vocab_size * sizeof(char *)); for (uint32_t i = 0; i < tokenizer->vocab_size; i++) { freadCheck(&length, sizeof(unsigned char), 1, file); assert(length > 0); // every token should be at least one character char *token_bytes = (char *)mallocCheck(length + 1); freadCheck(token_bytes, sizeof(char), length, file); token_bytes[length] = '\0'; // Add null terminator for printing tokenizer->token_table[i] = token_bytes; } // cleanups fcloseCheck(file); tokenizer->init_ok = 1; } const char *tokenizer_decode(Tokenizer *tokenizer, uint32_t token_id) { if (tokenizer->init_ok == 0) { return NULL; } if (token_id < tokenizer->vocab_size) { return tokenizer->token_table[token_id]; } else { printf("invalid token id %d!\n", token_id); return NULL; } } void tokenizer_free(Tokenizer *tokenizer) { if (tokenizer->init_ok) { for (uint32_t i = 0; i < tokenizer->vocab_size; i++) { free(tokenizer->token_table[i]); } free(tokenizer->token_table); } } // ---------------------------------------------------------------------------- // Logger lite, will probably grow/change some over time typedef struct { FILE *logfile; int flush_every; // every how many steps to flush the log } 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, poor man's argparse void error_usage() { // default run = debugging run with TinyShakespeare // bigger run = train on TinyStories! e.g. val/sample less often, but sample more tokens, write to logfile fprintf(stderr, "Usage: ./train_gpt2cu [options]\n"); fprintf(stderr, "Example: ./train_gpt2cu -i data/TinyStories -v 100 -s 100 -g 144 -o stories.log\n"); fprintf(stderr, "Options:\n"); fprintf(stderr, " -i input dataset prefix (default = data/tiny_shakespeare)\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, how often we evaluate val loss (default = 20)\n"); fprintf(stderr, " -m val_max_batches, up to how many val batches to estimate val loss? (default = 20)\n"); fprintf(stderr, " -s sample_every, how often we inference the model (default = 20)\n"); fprintf(stderr, " -g genT, how many steps of inference we do (default = 64)\n"); exit(EXIT_FAILURE); } // ---------------------------------------------------------------------------- // main training loop int main(int argc, char *argv[]) { multi_gpu_config = multi_gpu_config_init(&argc, &argv); // read in the (optional) command line arguments const char* input_dataset_prefix = "data/tiny_shakespeare"; // or e.g. data/TinyStories const char* output_log_file = NULL; int B = 4; // batch size int T = 1024; // sequence length max float learning_rate = 3e-4f; int val_loss_every = 20; // every how many steps do we eval validation loss? int val_max_batches = 20; // how many batches max do we eval for validation loss? int sample_every = 20; // every how many steps to do inference? int genT = 64; // number of steps of inference we will do for (int i = 1; i < argc; i+=2) { if (i + 1 >= argc) { error_usage(); } // must have arg after flag if (argv[i][0] != '-') { error_usage(); } // must start with dash if (strlen(argv[i]) != 2) { error_usage(); } // must be -x (one dash, one letter) // 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]); } // 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]); } else if (argv[i][1] == 'm') { val_max_batches = 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(); } } printf0("+-----------------------+----------------------------------------------------+\n"); printf0("| Parameter | Value |\n"); printf0("+-----------------------+----------------------------------------------------+\n"); printf0("| input dataset prefix | %-50s |\n", input_dataset_prefix); printf0("| output log file | %-50s |\n", output_log_file == NULL ? "NULL" : output_log_file); printf0("| batch size B | %-50d |\n", B); printf0("| sequence length T | %-50d |\n", T); printf0("| learning rate | %-50f |\n", learning_rate); printf0("| val_loss_every | %-50d |\n", val_loss_every); printf0("| val_max_batches | %-50d |\n", val_max_batches); printf0("| sample_every | %-50d |\n", sample_every); printf0("| genT | %-50d |\n", genT); printf0("+-----------------------+----------------------------------------------------+\n"); // set up the device cudaCheck(cudaSetDevice(multi_gpu_config.local_device_idx)); cudaDeviceProp deviceProp; cudaGetDeviceProperties(&deviceProp, multi_gpu_config.local_device_idx); // setup cuBLAS and cuBLASLt cublasCheck(cublasCreate(&cublas_handle)); cublasCheck(cublasLtCreate(&cublaslt_handle)); // TF32 precision is equivalent to torch.set_float32_matmul_precision('high') 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)); cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); printf0("| device | %-50s |\n", deviceProp.name); printf0("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); printf0("| precision | %-50s |\n", precision_mode_str); printf0("+-----------------------+----------------------------------------------------+\n"); // build the GPT-2 model from a checkpoint GPT2 model; gpt2_build_from_checkpoint(&model, load_filename); printf0("| load_filename | %-50s |\n", load_filename); printf0("| max_sequence_length T | %-50d |\n", model.config.max_seq_len); printf0("| vocab_size V | %-50d |\n", model.config.vocab_size); printf0("| num_layers L | %-50d |\n", model.config.num_layers); printf0("| num_heads NH | %-50d |\n", model.config.num_heads); printf0("| channels C | %-50d |\n", model.config.channels); printf0("| num_parameters | %-50zu |\n", model.num_parameters); printf0("+-----------------------+----------------------------------------------------+\n"); // build DataLoaders for both train and val char train_tokens_filename[128]; char val_tokens_filename[128]; assert(strlen(input_dataset_prefix) < 100); // being bit lazy here, make sure we don't overflow 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, &multi_gpu_config, train_tokens_filename, B, T); DataLoader val_loader; 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; printf0("| train_num_batches | %-50d |\n", train_num_batches); printf0("| val_num_batches | %-50d |\n", val_num_batches); printf0("+-----------------------+----------------------------------------------------+\n"); // pretty print in a table the multi-gpu configuration as well printf0("| num_processes | %-50d |\n", multi_gpu_config.num_processes); printf0("+-----------------------+----------------------------------------------------+\n"); // more prints related to allocations from gpt2_build_from_checkpoint down here to not mess up our table above printf0("num_parameters: %zu ==> bytes: %zu\n", model.num_parameters, model.num_parameters_bytes); printf0("allocated %d MiB for model parameters\n", (int)round(model.num_parameters_bytes / (1024 * 1024))); // set up the Logger Logger logger; logger_init(&logger, output_log_file); // build the Tokenizer Tokenizer tokenizer; tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); // some memory for generating samples from the model unsigned long long rng_state = 1337; int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); floatX* cpu_logits_raw = (floatX*)mallocCheck(model.config.vocab_size * sizeof(floatX)); float* cpu_logits = (float*)mallocCheck(model.config.vocab_size * sizeof(float)); // train 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; // once in a while estimate the validation loss 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; 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 (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; } // now sample from the model autoregressively printf("generating:\n---\n"); for (int t = 1; t < genT; t++) { // note that inference is very wasteful here because for each token // we re-calculate the forward pass for all of (B,T) positions from scratch // but the inference here is just for sanity checking anyway // and we can maybe optimize a bit more later, with careful tests gpt2_forward(&model, gen_tokens, NULL, B, T); // furthermore, below we're only using b=0 (i.e. the first row) of all B rows // we're in principle running B "inference streams" in parallel here // only using position 0 because it's a bit faster (copy less probs from GPU -> CPU) // get the V-dimensional vector probs[0, t-1, :] floatX* logits = model.acts.output + (t - 1) * model.config.vocab_size; // move probs back to CPU and sample cudaCheck(cudaMemcpy(cpu_logits_raw, logits, model.config.vocab_size * sizeof(floatX), cudaMemcpyDeviceToHost)); // convert to FP32 into cpu_logits (this does nothing useful if floatX == float) for (int i = 0; i < model.config.vocab_size; i++) { cpu_logits[i] = (float)cpu_logits_raw[i]; } float coin = random_f32(&rng_state); int next_token = sample_softmax(cpu_logits, model.config.vocab_size, coin); gen_tokens[t] = next_token; // print the generated token, either using the Tokenizer or a fallback if (tokenizer.init_ok) { const char* token_str = tokenizer_decode(&tokenizer, next_token); safe_printf(token_str); } else { // fall back to printing the token id printf("%d ", next_token); } fflush(stdout); } printf("\n---\n"); } // bit confusing: we want to make sure to eval and sample on 0th iteration // but also after the very last iteration. so we loop for step <= train_num_batches // instead of just < train_num_batches (one extra due to <=), only to do // the validation/sampling one last time, and then we break right here as we're done. if (last_step) { break; } // do a 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_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 = multi_gpu_config.num_processes * (B * T) / time_elapsed_s; printf0("step %4d/%d: train loss %f (acc %f) (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, model.accumulated_mean_loss, time_elapsed_s * 1000, tokens_per_second); logger_log_train(&logger, step, model.mean_loss); } // add a total average, for optimizations that are only mild improvements printf0("total average iteration time: %f ms\n", total_sum_iteration_time_s / train_num_batches * 1000); // free dataloader_free(&train_loader); dataloader_free(&val_loader); tokenizer_free(&tokenizer); gpt2_free(&model); free(cpu_logits_raw); free(cpu_logits); free(gen_tokens); cudaCheck(cudaFree(cublaslt_workspace)); cublasCheck(cublasDestroy(cublas_handle)); cublasCheck(cublasLtDestroy(cublaslt_handle)); logger_free(&logger); multi_gpu_config_free(&multi_gpu_config); return 0; } #endif