diff --git a/.github/workflows/ci_gpu.yml b/.github/workflows/ci_gpu.yml index d151506..9162d3b 100644 --- a/.github/workflows/ci_gpu.yml +++ b/.github/workflows/ci_gpu.yml @@ -31,19 +31,19 @@ jobs: run: python train_gpt2.py - name: Compile training and testing program - run: make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu + run: make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu - name: Train model (With OpenMP) run: OMP_NUM_THREADS=8 ./train_gpt2cu - name: Train model (FP32) with gpt2_124M.bin - run: | + run: | PRECISION=FP32 make train_gpt2cu ./train_gpt2cu -b 4 -t 64 -l 1e-4 -v 200 -s 200 -a 1 -x 10 -e gpt2_124M.bin - + - name: Build FP32 precision run: PRECISION=FP32 make test_gpt2cu profile_gpt2cu - + - name: Run default run: ./test_gpt2cu @@ -52,7 +52,7 @@ jobs: - name: Run recompute LN run: ./test_gpt2cu -r 2 - + - name: Build BF16 precision run: PRECISION=BF16 make train_gpt2cu test_gpt2cu profile_gpt2cu @@ -67,18 +67,18 @@ jobs: - name: Run recompute LN run: ./test_gpt2cu -r 2 - + - name: Train model fp32 (With OpenMP) run: OMP_NUM_THREADS=8 ./train_gpt2fp32cu - name: Execute testing program (With OpenMP) run: OMP_NUM_THREADS=8 ./test_gpt2cu - + - name: Execute testing program fp32 (With OpenMP) run: OMP_NUM_THREADS=8 ./test_gpt2fp32cu - name: Compile training and testing program without OpenMP - run: NO_OMP=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu + run: NO_OMP=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu - name: Train model (No OpenMP) run: NO_OMP=1 ./train_gpt2cu @@ -88,14 +88,14 @@ jobs: - name: Execute testing program (No OpenMP) run: ./test_gpt2cu -b 32 - + - name: Execute testing program fp32 (No OpenMP) run: ./test_gpt2fp32cu - name: Install cuDNN-frontend - run: + run: git clone https://github.com/NVIDIA/cudnn-frontend.git - + - name: Build with cuDNN run: USE_CUDNN=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu diff --git a/dev/cuda/adamw.cu b/dev/cuda/adamw.cu index 20a6560..74dfc2e 100644 --- a/dev/cuda/adamw.cu +++ b/dev/cuda/adamw.cu @@ -159,7 +159,7 @@ int main(int argc, char **argv) { // create random data on host (to be used for the CPU reference implementation) float* params_memory = make_random_float(num_parameters); float* grads_memory = make_random_float(num_parameters); - float* m_memory = make_random_float_01(num_parameters); + float* m_memory = make_random_float(num_parameters); float* v_memory = make_random_float_01(num_parameters); // move to GPU diff --git a/dev/cuda/attention_backward.cu b/dev/cuda/attention_backward.cu index c97dbee..f6d258d 100644 --- a/dev/cuda/attention_backward.cu +++ b/dev/cuda/attention_backward.cu @@ -68,7 +68,7 @@ void attention_forward_cpu(float* out, float* preatt, float* att, float* att_bth = att + b*NH*T*T + h*T*T + t*T; // pass 1: calculate query dot key and maxval - float maxval = -10000.0f; // TODO something better + float maxval = -FLT_MAX; for (int t2 = 0; t2 < T; t2++) { // used to be t2 <= t float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key diff --git a/dev/cuda/attention_forward.cu b/dev/cuda/attention_forward.cu index b632b4a..66b6a1b 100644 --- a/dev/cuda/attention_forward.cu +++ b/dev/cuda/attention_forward.cu @@ -98,7 +98,7 @@ void attention_forward_cpu(float* out, float* preatt, float* att, float* att_bth = att + b*NH*T*T + h*T*T + t*T; // pass 1: calculate query dot key and maxval - float maxval = -10000.0f; // TODO something better + float maxval = -FLT_MAX; for (int t2 = 0; t2 <= t; t2++) { const float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key @@ -203,7 +203,7 @@ __global__ void attention_softmax_kernel1(float* att, const float* preatt, float* att_bth = att + b*NH*T*T + h*T*T + t*T; // find maxval - float maxval = -10000.0f; // TODO something better + float maxval = -FLT_MAX; for (int t2 = 0; t2 <= t; t2++) { if (preatt_bth[t2] > maxval) { maxval = preatt_bth[t2]; diff --git a/dev/cuda/classifier_fused.cu b/dev/cuda/classifier_fused.cu index 5b6c986..9d93f9d 100644 --- a/dev/cuda/classifier_fused.cu +++ b/dev/cuda/classifier_fused.cu @@ -114,7 +114,7 @@ __device__ SoftmaxParams prepare_softmax(cg::thread_block_tile<32>& warp, int64_t idx, const float* inp, int V, int P) { // this warp (of 32) threads processes one row of inp, i.e. inp[idx, :] of shape (V,) // note that inp is actually (B * T, P) but we only use the first V elements - // this function tehen calculates: + // this function then calculates: // 1) the max value to subtract for numerical stability and // 2) the sum normalization factor const float* x = inp + idx * P; @@ -481,33 +481,6 @@ __global__ void fused_classifier_kernel4(floatX* dlogits, floatX* losses, floatX } } -// todo - move to common.h - or ideally somewhere it's not duplicated between train & common? -// requires all 32 threads in the warp to be active, but should work for any block size -// uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes -// the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end -// but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1 -using reduction_func_t = float (*) (float); -template -__device__ float blockReduce(float val, bool final_sync=false, float out_of_bounds=0.0f) { - // two reductions of up to 1024 threads: - // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) - __shared__ float shared_val[32]; - const int lane_id = threadIdx.x % 32; - const int warp_id = threadIdx.x / 32; - const int num_warps = blockDim.x / 32; - - float warp_val = warp_reduction(val); - if (lane_id == 0) { shared_val[warp_id] = warp_val; } - __syncthreads(); - warp_val = (lane_id < num_warps) ? shared_val[lane_id] : out_of_bounds; - float block_val = warp_reduction(warp_val); - - if (final_sync) { - __syncthreads(); // only needed in loops when effectively reusing shared memory etc. - } - return block_val; -} - __device__ SoftmaxParams prepare_softmax_blockwide3(int64_t idx, const floatX* inp, int V, int P) { // same but not float4 // one row of inp, i.e. inp[idx, :] of shape (V,) @@ -707,8 +680,8 @@ int main(int argc, char **argv) { cudaCheck(cudaSetDevice(deviceIdx)); // create host memory of random numbers - float* logits = make_random_float_01(B * T * V); - float* probs = (float*)malloc(B * T * V * sizeof(float)); + float* logits = make_random_float(B * T * V); + float* probs = make_random_float_01(B * T * V); float* dlogits = (float*)malloc(B * T * V * sizeof(float)); float* losses = (float*)malloc(B * T * sizeof(float)); float* dlosses = make_random_float(B * T); @@ -787,6 +760,7 @@ int main(int argc, char **argv) { free(losses); free(dlosses); free(targets); + free(outliers); cudaCheck(cudaFree(d_dlogits)); cudaCheck(cudaFree(d_losses)); cudaCheck(cudaFree(d_logits)); diff --git a/dev/cuda/common.h b/dev/cuda/common.h index 6502baa..61a783a 100644 --- a/dev/cuda/common.h +++ b/dev/cuda/common.h @@ -5,6 +5,8 @@ #include #include +#define WARP_SIZE 32U +extern cudaDeviceProp deviceProp; template __host__ __device__ T ceil_div(T dividend, T divisor) { @@ -18,6 +20,39 @@ __device__ float warpReduceSum(float val) { return val; } +// requires all 32 threads in the warp to be active, but should work for any block size +// uses non-dynamic shared memory so every call increases shared memory requirements by 128 bytes +// the fact it's unique shared memory allows us to avoid an extra __syncthreads() call at the end +// but if called inside a loop, the shared memory will be implicitly reused, so set final_sync to 1 +using reduction_func_t = float (*) (float); + +template +__device__ inline float blockReduce(float val, bool final_sync, float out_of_bounds) { + // two reductions of up to 1024 threads: + // 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle) + __shared__ float shared_val[WARP_SIZE]; + const int lane_id = threadIdx.x % WARP_SIZE; + const int warp_id = threadIdx.x / WARP_SIZE; + const int num_warps = blockDim.x / WARP_SIZE; + + float warp_val = warp_reduction(val); + if (lane_id == 0) { shared_val[warp_id] = warp_val; } + __syncthreads(); + warp_val = (lane_id < num_warps) ? shared_val[lane_id] : out_of_bounds; + float block_val = warp_reduction(warp_val); + + if (final_sync) { + __syncthreads(); // only needed in loops when effectively reusing shared memory etc. + } + return block_val; +} + +// Helper function to call blockReduce with default arguments +template +__device__ inline float blockReduce(float val) { + return blockReduce(val, false, 0.0f); +} + // ---------------------------------------------------------------------------- // checking utils diff --git a/dev/cuda/crossentropy_softmax_backward.cu b/dev/cuda/crossentropy_softmax_backward.cu index 27521bf..65c72a2 100644 --- a/dev/cuda/crossentropy_softmax_backward.cu +++ b/dev/cuda/crossentropy_softmax_backward.cu @@ -99,7 +99,7 @@ int main(int argc, char **argv) { cudaCheck(cudaSetDevice(deviceIdx)); // create host memory of random numbers - float* probs = make_random_float(B * T * V); + float* probs = make_random_float_01(B * T * V); int* targets = make_random_int(B * T, V); float* dlosses = make_random_float(B * T); float* dlogits = make_zeros_float(B * T * V); diff --git a/dev/cuda/global_norm.cu b/dev/cuda/global_norm.cu index 6c2ed03..f54a35a 100644 --- a/dev/cuda/global_norm.cu +++ b/dev/cuda/global_norm.cu @@ -16,6 +16,7 @@ nvcc -O3 --use_fast_math global_norm.cu -o global_norm #define ENABLE_BF16 #include "common.h" +cudaDeviceProp deviceProp; float global_norm_cpu(const float* data, size_t count) { // accumulate in double so we have an accurate numerical reference @@ -89,6 +90,54 @@ __global__ void norm_kernel2(float* out, const T* data, size_t count) { } } +template +__global__ void norm_kernel3(float* out, const T* data, size_t count) { + size_t index = blockIdx.x * blockDim.x + threadIdx.x; + size_t grid_width = blockDim.x * gridDim.x; + float accumulator = 0.f; + for(size_t i = index; i < count; i += grid_width) { + accumulator += (float)data[i] * (float)data[i]; + } + // block-level reduce + float block_sum = blockReduce(accumulator); + if(threadIdx.x == 0) { + atomicAdd(out, block_sum); + } +} + +// Same as kernel3 but without atomic adds -> this allows us to have determinism due to the +// non associativity of floating point operations. Roughly same performance as kernel3. +template +__global__ void norm_kernel4(float* out, const T* data, size_t count) { + size_t index = blockIdx.x * blockDim.x + threadIdx.x; + size_t grid_width = blockDim.x * gridDim.x; + float accumulator = 0.f; + for(size_t i = index; i < count; i += grid_width) { + accumulator += (float)data[i] * (float)data[i]; + } + // block-level reduce + float block_sum = blockReduce(accumulator); + // each block accumulates its partial sum to out[blockIdx.x] + // we want to avoid using atomic add here so we combine this kernel with the aggregate kernel call + // that sums up the partial block sums + if(threadIdx.x == 0) { + out[blockIdx.x] = block_sum; + } +} + +__global__ void global_norm_aggregate_kernel(float* out, size_t count) { + size_t index = threadIdx.x; + // grab block sums from the previous kernel, use 0. as the neutral sum element + float block_sum = (index < count) ? out[index] : 0.f; + float sum = blockReduce(block_sum); + if(threadIdx.x == 0) { + out[0] = sum; // out[0] ends up with the final norm squared + } +} + +// ---------------------------------------------------------------------------- +// kernel launchers + template void global_norm1(float* out, const T* values, size_t count, int block_size) { // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. @@ -111,17 +160,54 @@ void global_norm2(float* out, const T* values, size_t count, int block_size) { cudaCheck(cudaGetLastError()); } +template +void global_norm3(float* out, const T* values, size_t count, int block_size) { + // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. + // having one block less than possible is a tiny performance hit, having + // one block too many is catastrophic, since it only can start once all the other + // blocks finish. anyway, I think cuda_threads_per_SM should be a multiple of 512 + // on all gpus, so the division really is going to be exact. + const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + norm_kernel3<<>>(out, values, count); + cudaCheck(cudaGetLastError()); +} + +template +void global_norm4(float* out, const T* values, size_t count, int block_size) { + if (block_size <= 64) { + block_size = 128; // to avoid triggering the assert below + } + // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. + // having one block less than possible is a tiny performance hit, having + // one block too many is catastrophic, since it only can start once all the other + // blocks finish. anyway, I think cuda_threads_per_SM should be a multiple of 512 + // on all gpus, so the division really is going to be exact. + const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; + assert(grid_size > 0); // gives a better error than letting the call below fail + assert(grid_size < 1024); // we want to later accumulate the block sums in a single block + norm_kernel4<<>>(out, values, count); + cudaCheck(cudaGetLastError()); + global_norm_aggregate_kernel<<<1, 1024>>>(out, grid_size); + cudaCheck(cudaGetLastError()); +} + void global_norm(int kernel_num, float* out, const floatX* values, size_t count, int block_size) { switch (kernel_num) { case 1: return global_norm1(out, values, count, block_size); case 2: return global_norm2(out, values, count, block_size); + case 3: + return global_norm3(out, values, count, block_size); + case 4: + return global_norm4(out, values, count, block_size); } } int main(int argc, const char **argv) { setup_main(); + cudaGetDeviceProperties(&deviceProp, 0); int C = 768; int L = 12; @@ -148,7 +234,7 @@ int main(int argc, const char **argv) { // move to GPU float* d_out; floatX* d_inp; - cudaCheck(cudaMalloc(&d_out, sizeof(float))); + cudaCheck(cudaMalloc(&d_out, 1024 * sizeof(float))); // 1024 needed for kernel 4 cudaCheck(cudaMalloc(&d_inp, num_params * sizeof(floatX))); cudaCheck(memcpy_convert(d_inp, inp, num_params)); diff --git a/dev/cuda/layernorm_backward.cu b/dev/cuda/layernorm_backward.cu index dc9d7e9..3930cec 100644 --- a/dev/cuda/layernorm_backward.cu +++ b/dev/cuda/layernorm_backward.cu @@ -874,7 +874,6 @@ __global__ void layernorm_backward_kernel9(floatX* dinp, floatX* dweight, floatX } __trap(); // prefer to crash here than run into a deadlock later on } - constexpr int WARP_SIZE = 32; int BLOCK_SIZE = blockDim.x; int warpsInBlock = BLOCK_SIZE / WARP_SIZE; //number of warps in block extern __shared__ float shared[]; // size = 2 * C + 1 @@ -1059,7 +1058,6 @@ layernorm_backward_kernel10(floatX* dinp, floatX* dweight, floatX* dbias, float* const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd, int B, int T, int C) { - constexpr int WARP_SIZE = 32; int BLOCK_SIZE = blockDim.x; int warpsInBlock = BLOCK_SIZE / WARP_SIZE; //number of warps in block extern __shared__ float shared[]; // size = 2 * C + 1 diff --git a/dev/cuda/layernorm_forward.cu b/dev/cuda/layernorm_forward.cu index 3e94828..8631aa1 100644 --- a/dev/cuda/layernorm_forward.cu +++ b/dev/cuda/layernorm_forward.cu @@ -290,7 +290,7 @@ __global__ void layernorm_forward_kernel5(float* __restrict__ out, float* __rest int num_warps = blockDim.x / 32; int warp_id = threadIdx.x / 32; int lane_id = threadIdx.x % 32; - int idx = blockIdx.x; // simpoy one block per row + int idx = blockIdx.x; // simply one block per row // the row of input that this group of threads is responsible for const float* x = inp + idx * C; // thread coarsening through the row, reduce the sum in series @@ -356,9 +356,9 @@ void layernorm_forward2(float* out, float* mean, float* rstd, const int block_size) { int N = B * T; // in mean and rstd, threads cooperate within blocks via reductions - mean_kernel<<>>(mean, inp, N, C, block_size); + mean_kernel<<>>(mean, inp, N, C, block_size); cudaCheck(cudaGetLastError()); - rstd_kernel<<>>(rstd, inp, mean, N, C, block_size); + rstd_kernel<<>>(rstd, inp, mean, N, C, block_size); cudaCheck(cudaGetLastError()); // in the normalization, everything just gets flattened out const int block_size2 = 256; @@ -394,6 +394,7 @@ void layernorm_forward5(float* out, float* mean, float* rstd, int B, int T, int C, const int block_size) { assert(block_size % 32 == 0); + assert(block_size <= 1024); const int N = B * T; const int grid_size = N; layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); @@ -473,9 +474,6 @@ int main(int argc, char **argv) { printf("Using kernel %d\n", kernel_num); int block_sizes[] = {32, 64, 128, 256, 512, 1024}; - float* out_gpu = (float*)malloc(B * T * C * sizeof(float)); - float* mean_gpu = (float*)malloc(B * T * sizeof(float)); - float* rstd_gpu = (float*)malloc(B * T * sizeof(float)); layernorm_forward_cpu(out, mean, rstd, inp, weight, bias, B, T, C); diff --git a/dev/cuda/matmul_backward_bias.cu b/dev/cuda/matmul_backward_bias.cu index ad56a8b..86fd373 100644 --- a/dev/cuda/matmul_backward_bias.cu +++ b/dev/cuda/matmul_backward_bias.cu @@ -2,7 +2,7 @@ Kernels for matmul backward pass bias only. Compile example: -nvcc -O3 -lcublas -lcublasLt matmul_backward_bias.cu -lineinfo -o matmul_backward_bias +nvcc -O3 -lcublas -lcublasLt -std=c++17 matmul_backward_bias.cu -lineinfo -o matmul_backward_bias ./matmul_backward_bias 1 ./matmul_backward_bias 2 @@ -116,7 +116,7 @@ __global__ void matmul_backward_bias_kernel2(floatX* dbias, const floatX* dout, sum = cg::reduce(warp, sum, cg::plus{}); // write the result to output (global memory) if(warp.thread_rank() == 0) { - dbias[idx] += (floatX)sum; + dbias[idx] = (float)dbias[idx] + sum; } } @@ -149,7 +149,7 @@ __global__ void matmul_backward_bias_kernel3(floatX* dbias, const floatX* dout, float block_sum = cg::reduce(warp, warp_sum, cg::plus{}); // sum(x) // write the result to output (global memory) if(threadIdx.x == 0) { - dbias[idx] += block_sum; + dbias[idx] = (float)dbias[idx] + block_sum; } } @@ -189,7 +189,7 @@ __global__ void matmul_backward_bias_kernel4(floatX* dbias, const floatX* dout, for (int j = 0; j < vstep; j++) { dout_sum += smem[lane_id + j * warpSize]; } - dbias[tl + lane_id] += dout_sum; + dbias[tl + lane_id] = (float)dbias[tl + lane_id] + dout_sum; } } diff --git a/llmc/adamw.cuh b/llmc/adamw.cuh index 94d4642..f806eaa 100644 --- a/llmc/adamw.cuh +++ b/llmc/adamw.cuh @@ -40,9 +40,7 @@ __device__ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* float param = old_param - (learning_rate * (m / (sqrtf(v) + eps) + weight_decay * old_param)); // update our low precision version of the parameters using stochastic rounding // this will be used in the next forward pass - // TODO: simply doing `params_memory[i] = (floatX)param;` breaks everything (why?) - unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x + blockDim.y * gridDim.y, seed); - stochastic_rounding(param, ¶ms_memory[idx], random); + stochastic_rounding(param, ¶ms_memory[idx], seed); // write the full, float version of the param into our master copy, if we maintain one // this will be used in the next update if (master_params_memory != NULL) { master_params_memory[idx] = param; } diff --git a/llmc/cuda_utils.cuh b/llmc/cuda_utils.cuh index 1bca60a..fe3b265 100644 --- a/llmc/cuda_utils.cuh +++ b/llmc/cuda_utils.cuh @@ -190,7 +190,8 @@ __device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY // 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); + // makes sure each thread gets a different random number + unsigned int random = Get2dNoiseUint(threadIdx.x, blockIdx.x * blockDim.x + blockIdx.y, seed); unsigned int threshold = random & 0xFFFF; unsigned int float_bits = __float_as_uint(in); unsigned int rounded_bits = float_bits & 0x0000FFFF; diff --git a/llmc/fused_classifier.cuh b/llmc/fused_classifier.cuh index b13fa35..a52765d 100644 --- a/llmc/fused_classifier.cuh +++ b/llmc/fused_classifier.cuh @@ -85,6 +85,12 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS) losses[idx] = (floatX)(-logf(prob)); } + // without this synchronization point we have a race condition: + // the logits used above to compute the loss are concurrently (race) modified to carry backward pass grads. + // since the "logits" are overwritten to be in the [-1, 1] range and sp.Offset is sometimes smaller than -90 + // we errouneously end up computing exp^(90+) which gives us infinities in the loss! this is the fix. + __syncthreads(); + // calculate the gradients directly, saves bandwidth from probs during training // but also supports writing probs for inference-only and debugging const floatX* logits_vec = logits + idx * P; diff --git a/llmc/gelu.cuh b/llmc/gelu.cuh index ce9aa12..cd5c297 100644 --- a/llmc/gelu.cuh +++ b/llmc/gelu.cuh @@ -50,7 +50,7 @@ __global__ void gelu_backward_inplace_kernel(floatX* d_in_out, const floatX* inp void gelu_forward(floatX* out, const floatX* inp, int N, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 512; - assert(N % block_size == 0); + assert(N % (block_size * x128::size) == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); gelu_forward_kernel2<<>>(out, inp); cudaCheck(cudaGetLastError()); @@ -59,7 +59,7 @@ void gelu_forward(floatX* out, const floatX* inp, int N, cudaStream_t stream) { void gelu_backward_inplace(floatX* d_in_out, const floatX* inp, const int N, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 128; - assert(N % block_size == 0); + assert(N % (block_size * x128::size) == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); gelu_backward_inplace_kernel<<>>(d_in_out, inp); cudaCheck(cudaGetLastError()); diff --git a/llmc/global_norm.cuh b/llmc/global_norm.cuh index 92866e8..9e23744 100644 --- a/llmc/global_norm.cuh +++ b/llmc/global_norm.cuh @@ -13,13 +13,7 @@ Global norm, used in gradient clipping template __device__ float global_norm_squared_for_range(const T* data, size_t count) { - // we want as few atomics as possible, so each block tries to do - // the maximum amount of work (so no fixed chunk, but instead iterating - // until we run out of data), and then we reduce inside the block - // and finally have just one atomic per block. - // out will be updated atomically from all thread blocks. It is a float, so the - // atomic op is unproblematic - size_t index = threadIdx.x + blockDim.x * blockIdx.x; + size_t index = blockIdx.x * blockDim.x + threadIdx.x; size_t grid_width = blockDim.x * gridDim.x; float accumulator = 0.f; for(size_t i = index; i < count; i += grid_width) { @@ -32,16 +26,47 @@ __device__ float global_norm_squared_for_range(const T* data, size_t count) { template __global__ void global_norm_squared_kernel(float* out, const T* data, size_t count, ptrdiff_t stride) { float block_sum = global_norm_squared_for_range(data + blockIdx.y * stride, count); + // each block accumulates its partial sum to out[out_index] + // we want to avoid using atomic add here so we combine this kernel with another kernel call + // that sums up the partial block sums if(threadIdx.x == 0) { - atomicAdd(out, block_sum); + size_t out_index = blockIdx.y * gridDim.x + blockIdx.x; + out[out_index] = out[out_index] + block_sum; + } +} + +__global__ void global_norm_aggregate_kernel(float* out, size_t grid_size) { + size_t index = threadIdx.x; + // grab block sums from the previous kernel, use 0. as the neutral sum element + float block_sum = (index < grid_size) ? out[index] : 0.f; + float sum = blockReduce(block_sum); + if(threadIdx.x == 0) { + out[0] = sum; // out[0] ends up with the final norm squared } } // ---------------------------------------------------------------------------- // kernel launcher +// Helper function determines the maximum number of block sums +int get_max_num_block_sums(int* num_slices_all, int numel) { + // NOTE: this needs to be kept in sync with `global_norm_squared` below. + const int block_size = 512; + const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; + assert(grid_size > 0); + int max_num_block_sums = 0; + for (int i = 0; i < numel; i++) { + int num_slices = num_slices_all[i]; + const int gx = CEIL_DIV(grid_size, num_slices); + const int gy = num_slices; + max_num_block_sums = max(max_num_block_sums, gx * gy); + } + + return max_num_block_sums; +} + template -void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t stride, int num_slices, bool reset, cudaStream_t stream) { +void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t stride, int num_slices, int max_num_block_sums, bool reset, cudaStream_t stream) { const int block_size = 512; // launch just enough blocks to fill the grid. deliberately no DIV_CEIL. // having one block less than possible is a tiny performance hit, having @@ -50,13 +75,22 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st // on all gpus, so the division really is going to be exact. const int grid_size = deviceProp.maxThreadsPerMultiProcessor * deviceProp.multiProcessorCount / block_size; assert(grid_size > 0); // gives a better error than letting the call below fail - // initialize out with zero - if(reset) { - cudaCheck(cudaMemsetAsync(out, 0, sizeof(float), stream)); - } + const int gx = CEIL_DIV(grid_size, num_slices); const int gy = num_slices; + + assert(gx * gy < 1024); // we want to later accumulate the block sums in a single block + + if (reset) { + cudaCheck(cudaMemsetAsync(out, 0, max_num_block_sums * sizeof(float), stream)); + } global_norm_squared_kernel<<>>(out, values, count, stride); cudaCheck(cudaGetLastError()); } +void global_norm_squared_aggregate(float* out, int max_num_block_sums, cudaStream_t stream) { + assert(max_num_block_sums > 0 && max_num_block_sums < 1024); // we need to accumulate the block sums in a single block + // important to use 1024 here for determinism, otherwise blockreduce might introduce errors + global_norm_aggregate_kernel<<<1, 1024, 0, stream>>>(out, max_num_block_sums); + cudaCheck(cudaGetLastError()); +} diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index a49f0bf..9c9afa3 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -368,7 +368,7 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, void residual_forward(floatX* out, const floatX* inp1, const floatX* inp2, int N, cudaStream_t stream) { NVTX_RANGE_FN(); const int block_size = 256; - assert(N % block_size == 0); + assert(N % (block_size * x128::size) == 0); const int grid_size = CEIL_DIV(N, block_size * x128::size); residual_forward_kernel<<>>(out, inp1, inp2); cudaCheck(cudaGetLastError()); diff --git a/requirements.txt b/requirements.txt index 80471a8..ea4bc76 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ tqdm -numpy +numpy<2 torch tiktoken transformers diff --git a/train_gpt2.cu b/train_gpt2.cu index 878f0c9..1045b6f 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -970,32 +970,35 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl float* grad_norm_squared = (float*)model->acts.output; float grad_norm_squared_cpu = 0.0f; + int num_slices[2] = {1, model->config.num_layers}; + int max_num_block_sums = get_max_num_block_sums(num_slices, 2); if (multi_gpu_config->zero_stage == 1) { // because of the ncclReduceScatter() in backward, // grads_memory only contains the averaged gradients at the local shards, // so we only calculate the grad norm at the grads_memory belonging to the local shards - cudaCheck(cudaMemsetAsync(grad_norm_squared, 0, sizeof(float), main_stream)); for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { if((i < 2 || i > 13)) { ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i); ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); ptrdiff_t offset = tensor.offset + shard.offset; - global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, false, main_stream); + global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, 0, 1, max_num_block_sums, i == 0, main_stream); } else { ShardInfo tensor = gpt2_get_tensor_at_layer(model, 0, i); ShardInfo shard = multi_gpu_get_shard_offset(tensor.size, multi_gpu_config, 1); ptrdiff_t offset = tensor.offset + shard.offset; global_norm_squared(grad_norm_squared, grads_memory + offset, shard.size, tensor.size, model->config.num_layers, - false, main_stream); + max_num_block_sums, false, main_stream); } } + global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream); cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); // further sum the (partial) squared norm across all GPUs (see comment ^1 above) grad_norm_squared_cpu = multi_gpu_cpu_float_sum(grad_norm_squared_cpu); } else { // in regular DDP, backward has averaged the gradients across all GPUs // so each GPU can compute the squared norm over the whole grad vector, with no added comms needed - global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters, 0, 1, true, main_stream); + global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters, 0, 1, max_num_block_sums, true, main_stream); + global_norm_squared_aggregate(grad_norm_squared, max_num_block_sums, main_stream); cudaCheck(cudaMemcpy(&grad_norm_squared_cpu, grad_norm_squared, sizeof(float), cudaMemcpyDeviceToHost)); } @@ -1009,10 +1012,12 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl float grad_scale = (grad_norm_cpu > grad_clip) ? grad_clip / grad_norm_cpu : 1.0f; // AdamW update - unsigned int seed = random_u32(&model->rng_state); // handle adamw for all the transformer blocks for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) { + // generate a unique seed for each tensor + unsigned int seed = random_u32(&model->rng_state); + int num_layers = model->config.num_layers; if((i < 2 || i > 13)) { num_layers = 1; @@ -1037,7 +1042,7 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl float* master_ptr = NULL; if (model->master_weights != NULL) { master_ptr = model->master_weights + opt_state_offset; } if(init_master_weights) { - size_t grid_size = CEIL_DIV(shard_num_parameters, 512); + size_t grid_size = CEIL_DIV(shard.size, 512); copy_and_cast_kernel<<>>(master_ptr, param_ptr, shard.size, shard.size, tensor.size); cudaCheck(cudaGetLastError()); @@ -1173,7 +1178,8 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) state_header[1] = 1; // version number state_header[2] = multi_gpu_config.num_processes; // number of processes state_header[3] = multi_gpu_config.process_rank; // rank of this process - state_header[4] = loader->should_shuffle; // shuffle state of the dataloader + state_header[4] = model->use_master_weights; // whether we're using fp32 master weights + state_header[5] = loader->should_shuffle; // shuffle state of the dataloader // int main state, start at 10 to leave some padding state_header[10] = step; // step of the optimization // model rng state, start at 20 to leave some padding @@ -1190,8 +1196,13 @@ void save_state(const char* filename, int step, GPT2* model, DataLoader* loader) fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(cpu_buffer, model->v_memory, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + if (model->use_master_weights == 1) { + cudaCheck(cudaMemcpy(cpu_buffer, model->master_weights, shard_num_parameters * sizeof(float), cudaMemcpyDeviceToHost)); + fwrite(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + } free(cpu_buffer); + // write dataloader state if we are using the Permuted version of it if (loader->should_shuffle) { fwrite(&loader->glob_result.gl_pathc, sizeof(size_t), 1, state_file); // number of shards fwrite(loader->shard_indices, sizeof(int), loader->glob_result.gl_pathc, state_file); @@ -1210,49 +1221,62 @@ void load_state(int* step, GPT2* model, DataLoader* loader, const char* filename assert(state_header[1] == 1); // version number assert(state_header[2] == multi_gpu_config.num_processes); // number of processes assert(state_header[3] == multi_gpu_config.process_rank); // rank of this process - int should_shuffle = state_header[4]; // shuffle state of the dataloader + int use_master_weights = state_header[4]; // whether we're using fp32 master weights + int should_shuffle = state_header[5]; // shuffle state of the dataloader *step = state_header[10]; // step of the optimization model->rng_state = *((unsigned long long*)&state_header[20]); // random number generator state size_t current_shard_idx = *((size_t*)&state_header[30]); // shard index size_t current_sample_idx = *((size_t*)&state_header[32]); // position in shard - // read AdamW m, v (they are all float) - // also allocate the m, v memory in the model, if it does not yet exist + // read AdamW m, v, master_weights (they are all float) + // allocate all the needed memory as necessary size_t shard_num_parameters = multi_gpu_config.shard_num_parameters; if (model->m_memory == NULL) { printf0("allocating %zu MiB for AdamW optimizer state m\n", (shard_num_parameters * sizeof(float)) >> 20); - printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->m_memory, shard_num_parameters * sizeof(float))); + } + if (model->v_memory == NULL) { + printf0("allocating %zu MiB for AdamW optimizer state v\n", (shard_num_parameters * sizeof(float)) >> 20); cudaCheck(cudaMalloc((void**)&model->v_memory, shard_num_parameters * sizeof(float))); } + if (model->master_weights == NULL && use_master_weights == 1) { + printf0("allocating %zu MiB for master copy of params\n", (shard_num_parameters * sizeof(float)) >> 20); + cudaCheck(cudaMalloc((void**)&model->master_weights, shard_num_parameters * sizeof(float))); + } float* cpu_buffer = (float*)mallocCheck(shard_num_parameters * sizeof(float)); freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(model->m_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); cudaCheck(cudaMemcpy(model->v_memory, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); + if (use_master_weights == 1) { + freadCheck(cpu_buffer, sizeof(float), shard_num_parameters, state_file); + cudaCheck(cudaMemcpy(model->master_weights, cpu_buffer, shard_num_parameters * sizeof(float), cudaMemcpyHostToDevice)); + } free(cpu_buffer); - if (should_shuffle) { - loader->should_shuffle = 1; - + // revive the DataLoader object and its state + loader->should_shuffle = should_shuffle; + if (should_shuffle == 1) { + // ensure the number of shards matches size_t glob_result_gl_pathc; freadCheck(&glob_result_gl_pathc, sizeof(size_t), 1, state_file); assert(glob_result_gl_pathc == loader->glob_result.gl_pathc); - + // read the shard indices loader->shard_indices = (int*)mallocCheck(loader->glob_result.gl_pathc * sizeof(int)); freadCheck(loader->shard_indices, sizeof(int), loader->glob_result.gl_pathc, state_file); - + // ensure the number of samples matches size_t shard_num_samples; freadCheck(&shard_num_samples, sizeof(size_t), 1, state_file); assert(shard_num_samples == loader->shard_num_samples); - + // read the intra-shard indices loader->intra_shard_indices = (int*)mallocCheck(loader->shard_num_samples * sizeof(int)); freadCheck(loader->intra_shard_indices, sizeof(int), loader->shard_num_samples, state_file); - + // read the shuffle rng state freadCheck(&loader->shuffle_rng, sizeof(mt19937_state), 1, state_file); } - dataloader_resume(loader, current_shard_idx, current_sample_idx); + + // all done, close state file fclose(state_file); }