diff --git a/dev/cuda/fused_residual_forward.cu b/dev/cuda/fused_residual_forward.cu index b98a67c..9752873 100644 --- a/dev/cuda/fused_residual_forward.cu +++ b/dev/cuda/fused_residual_forward.cu @@ -133,7 +133,7 @@ __global__ void fused_residual_forward2(floatX* residual, floatX* normed, floatX for(int c = 0; c < C; ++c) { float out = (float)inp1[c] + (float)inp2[c]; m += out; - residual[c] = out; + residual[c] = (floatX)out; } m = m / C; @@ -149,11 +149,11 @@ __global__ void fused_residual_forward2(floatX* residual, floatX* normed, floatX for (int c = 0; c < C; c++) { float n = (s * ((float)residual[c] - m)); // normalized output float o = n * (float)weight[c] + (float)bias[c]; // scale and shift it - normed[c] = o; // write + normed[c] = (floatX)o; // write } // cache the mean and rstd for the backward pass later - mean[idx] = m; - rstd[idx] = s; + mean[idx] = (floatX)m; + rstd[idx] = (floatX)s; } // handle one token per warp for coalesced access @@ -232,7 +232,7 @@ __global__ void fused_residual_forward_kernel4(floatX* residual, floatX* normed, const x128 in2 = load128cs(inp2 + c); x128 out; for(int k = 0; k < x128::size; ++k) { - out[k] = (float)in1[k] + (float)in2[k]; + out[k] = (floatX)((float)in1[k] + (float)in2[k]); sum += (float)out[k]; sum_sq += (float)out[k] * (float)out[k]; } @@ -309,7 +309,7 @@ __global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, const x128 in2 = load128cs(inp2 + c); x128 out; for(int k = 0; k < x128::size; ++k) { - out[k] = (float)in1[k] + (float)in2[k]; + out[k] = (floatX)((float)in1[k] + (float)in2[k]); sum += (float)out[k]; } store128cs(residual + c, out); @@ -372,8 +372,8 @@ __global__ void fused_residual_forward_kernel6(floatX* residual, floatX* normed, // weights and biases are shared among all tokens x128* s_weight = reinterpret_cast(params); x128* s_bias = reinterpret_cast(params + C * sizeof(floatX)); - // residual output (input to layernorm) is indpendent for each sub-block indicates by threadIdx.z - x128* s_res = reinterpret_cast(params + (2 + threadIdx.z) * C * sizeof(floatX) ); + // residual output (input to layernorm) is independent for each sub-block indicates by threadIdx.z + x128* s_res = reinterpret_cast(params + (2 + threadIdx.z) * C * sizeof(floatX)); // similarly, each sub-block needs its own reduction buffers float* s_mean = reinterpret_cast(params + (2 + blockDim.z) * C * sizeof(floatX) + threadIdx.z * 32 * sizeof(float)); float* s_var = reinterpret_cast(params + (2 + blockDim.z) * C * sizeof(floatX) + 32 * sizeof(float) * (blockDim.z + threadIdx.z)); @@ -385,10 +385,10 @@ __global__ void fused_residual_forward_kernel6(floatX* residual, floatX* normed, s_weight[c / x128::size] = load128(weight + c); s_bias[c / x128::size] = load128(bias + c); } + // the block-level reductions will cause sync before the first time we read these // => no syncthreads needed here - // loop over all tokens for(int tidx = blockIdx.x * blockDim.z + threadIdx.z; tidx < N; tidx += gridDim.x * blockDim.z) { // adjust pointers to current token diff --git a/dev/cuda/layernorm_forward.cu b/dev/cuda/layernorm_forward.cu index 8631aa1..4e4552b 100644 --- a/dev/cuda/layernorm_forward.cu +++ b/dev/cuda/layernorm_forward.cu @@ -29,6 +29,8 @@ verstion 5 allocates blocks per row instead of warps per row, same alg as 4 othe #include #include "common.h" +#define WARP_SIZE 32 + // ---------------------------------------------------------------------------- // CPU code reference @@ -337,6 +339,82 @@ __global__ void layernorm_forward_kernel5(float* __restrict__ out, float* __rest } } +// Inspired by `fused_residual_forward_kernel5` in fused_residual_forward.cu +__global__ void layernorm_forward_kernel6(float* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd, + const float* __restrict__ inp, const float* __restrict__ weight, + const float* __restrict__ bias, int N, int C) { + assert(blockDim.x == WARP_SIZE); + + // load weights and biases into shared memory + // do this before we allow any threads to exit! + extern __shared__ char params[]; + // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so + // let's keep everything as x128 + x128* s_weight = reinterpret_cast(params); + x128* s_bias = reinterpret_cast(params) + (C / x128::size); + x128* s_in = reinterpret_cast(params) + ((2 + threadIdx.y) * C / x128::size); + + int sidx = (threadIdx.x + WARP_SIZE * threadIdx.y) * x128::size; + for(int i = sidx; i < C; i += blockDim.y * WARP_SIZE * x128::size) { + s_weight[i/x128::size] = load128(weight + i); + s_bias[i/x128::size] = load128(bias + i); + } + __syncthreads(); + + int idx = blockIdx.x * blockDim.y + threadIdx.y; + if(idx >= N) { return; } // guard + + // adjust pointers to current token + inp += idx * C; + out += idx * C; + + const float eps = 1e-5f; + float sum = 0.0f; + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = load128cs(inp + c); + for(int k = 0; k < x128::size; ++k) { + sum += (float)in_data[k]; + } + s_in[c / x128::size] = in_data; + } + + sum = warpReduceSum(sum); + float m = sum / C; + float v = 0.f; + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + for(int k = 0; k < x128::size; ++k) { + v += ((float)in_data[k] - m) * ((float)in_data[k] - m); + } + } + + v = warpReduceSum(v) / C; + float s = rsqrtf(v + eps); + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + const x128 w = s_weight[c / x128::size]; + const x128 b = s_bias[c / x128::size]; + x128 out_data; + for(int k = 0; k < x128::size; ++k) { + float n = s * ((float)in_data[k] - m); // normalized output + float o = n * (float)w[k] + (float)b[k]; // scale and shift it + out_data[k] = o; + } + + store128cs(out + c, out_data); + } + // cache the mean and rstd for the backward pass later + if(threadIdx.x == 0 && mean != nullptr) { + __stcs(mean + idx, m); + } + // store the rstd, no need to cache it + if(threadIdx.x == 0 && rstd != nullptr) { + __stcs(rstd + idx, s); + } +} + // ---------------------------------------------------------------------------- // kernel launcher @@ -401,6 +479,31 @@ void layernorm_forward5(float* out, float* mean, float* rstd, cudaCheck(cudaGetLastError()); } +void layernorm_forward6(float* out, float* mean, float* rstd, + const float* inp, const float* weight, const float* bias, + int B, int T, int C, + int block_size) { + assert(block_size % 32 == 0); + const int N = B * T; + int block_y = block_size / WARP_SIZE; + const int grid_size = ceil_div(N, block_y); + size_t smem = (2 + block_y) * C * sizeof(float); + + // in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute + // this may fail, in which case we fall back to the smem free implementation. + cudaCheck(cudaGetLastError()); + auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + cudaGetLastError(); + if (status == cudaSuccess) { + layernorm_forward_kernel6<<>>(out, mean, rstd, inp, weight, bias, N, C); + } else { + const int grid_size = N; + // fall back to the version without shared memory + layernorm_forward_kernel5<<>>(out, mean, rstd, inp, weight, bias, N, C); + } + cudaCheck(cudaGetLastError()); +} + // kernel version dispatch void layernorm_forward(int kernel_num, float* out, float* mean, float* rstd, @@ -423,6 +526,9 @@ void layernorm_forward(int kernel_num, case 5: layernorm_forward5(out, mean, rstd, inp, weight, bias, B, T, C, block_size); break; + case 6: + layernorm_forward6(out, mean, rstd, inp, weight, bias, B, T, C, block_size); + break; default: printf("Invalid kernel number\n"); exit(1); diff --git a/llmc/dataloader.h b/llmc/dataloader.h index d6a5b23..f3b0fca 100644 --- a/llmc/dataloader.h +++ b/llmc/dataloader.h @@ -110,7 +110,8 @@ void prepare_intra_shard_indices_(DataLoader *loader) { free(loader->intra_shard_indices); } loader->intra_shard_indices = (int*)mallocCheck(loader->shard_num_samples * sizeof(int)); - random_permutation_with_init(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng, 1); + init_identity_permutation(loader->intra_shard_indices, loader->shard_num_samples); + random_permutation(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng); } void dataloader_reset(DataLoader *loader) { @@ -118,7 +119,7 @@ void dataloader_reset(DataLoader *loader) { loader->current_sample_idx = 0; if (loader->should_shuffle) { // shuffle the shards - random_permutation_with_init(loader->shard_indices, loader->glob_result.gl_pathc, &loader->shuffle_rng, 0); + random_permutation(loader->shard_indices, loader->glob_result.gl_pathc, &loader->shuffle_rng); } dataloader_load_shard_(loader, loader->current_shard_idx); @@ -178,9 +179,7 @@ void dataloader_init(DataLoader *loader, manual_seed(&shuffle_rng, 42 + process_rank); loader->shuffle_rng = shuffle_rng; loader->shard_indices = (int*)mallocCheck(loader->glob_result.gl_pathc * sizeof(int)); - for (int i = 0; i < loader->glob_result.gl_pathc; i++) { - loader->shard_indices[i] = i; // start with identity permutation - } + init_identity_permutation(loader->shard_indices, loader->glob_result.gl_pathc); loader->intra_shard_indices = NULL; // dynamically allocated allowing different shard sizes } diff --git a/llmc/layernorm.cuh b/llmc/layernorm.cuh index 9c9afa3..681155e 100644 --- a/llmc/layernorm.cuh +++ b/llmc/layernorm.cuh @@ -64,6 +64,81 @@ __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __re } } +__global__ void layernorm_forward_kernel6(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd, + const floatX* __restrict__ inp, const floatX* __restrict__ weight, + const floatX* __restrict__ bias, int N, int C) { + assert(blockDim.x == WARP_SIZE); + + // load weights and biases into shared memory + // do this before we allow any threads to exit! + extern __shared__ char* params[]; + // load128/store128 sometimes generated multiple instructions when the types here were floatX*, so + // let's keep everything as x128 + x128* s_weight = reinterpret_cast(params); + x128* s_bias = reinterpret_cast(params) + (C / x128::size); + x128* s_in = reinterpret_cast(params) + ((2 + threadIdx.y) * C / x128::size); + + int sidx = (threadIdx.x + WARP_SIZE * threadIdx.y) * x128::size; + for(int i = sidx; i < C; i += blockDim.y * WARP_SIZE * x128::size) { + s_weight[i/x128::size] = load128(weight + i); + s_bias[i/x128::size] = load128(bias + i); + } + __syncthreads(); + + int idx = blockIdx.x * blockDim.y + threadIdx.y; + if(idx >= N) { return; } // guard + + // adjust pointers to current token + inp += idx * C; + out += idx * C; + + const float eps = 1e-5f; + float sum = 0.0f; + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = load128cs(inp + c); + for(int k = 0; k < x128::size; ++k) { + sum += (float)in_data[k]; + } + s_in[c / x128::size] = in_data; + } + + sum = warpReduceSum(sum); + float m = sum / C; + float v = 0.f; + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + for(int k = 0; k < x128::size; ++k) { + v += ((float)in_data[k] - m) * ((float)in_data[k] - m); + } + } + + v = warpReduceSum(v) / C; + float s = rsqrtf(v + eps); + + for(int c = threadIdx.x * x128::size; c < C; c += WARP_SIZE * x128::size) { + const x128 in_data = s_in[c / x128::size]; + const x128 w = s_weight[c / x128::size]; + const x128 b = s_bias[c / x128::size]; + x128 out_data; + for(int k = 0; k < x128::size; ++k) { + float n = s * ((float)in_data[k] - m); // normalized output + float o = n * (float)w[k] + (float)b[k]; // scale and shift it + out_data[k] = (floatX)o; + } + + store128cs(out + c, out_data); + } + // cache the mean and rstd for the backward pass later + if(threadIdx.x == 0 && mean != nullptr) { + __stcs(mean + idx, (floatX)m); + } + // store the rstd, no need to cache it + if(threadIdx.x == 0 && rstd != nullptr) { + __stcs(rstd + idx, (floatX)s); + } +} + __global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, floatX* mean, floatX* rstd, const floatX* inp1, const floatX* inp2, const floatX* weight, const floatX* bias, @@ -354,14 +429,29 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with // ---------------------------------------------------------------------------- // kernel launchers +// similar to `fused_residual_forward5` void layernorm_forward(floatX* out, floatX* mean, floatX* rstd, floatX* inp, const floatX* weight, const floatX* bias, int B, int T, int C, cudaStream_t stream) { NVTX_RANGE_FN(); - const int block_size = 512; + const int block_size = 256; + int block_y = block_size / WARP_SIZE; const int N = B * T; - const int grid_size = CEIL_DIV(N * WARP_SIZE, block_size); - layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); + const int grid_size = CEIL_DIV(N, block_y); + size_t smem = (2 + block_y) * C * sizeof(floatX); + + // in order to use more than 48 KiB of smem, need to call cudaFuncSetAttribute + // this may fail, in which case we fall back to the smem free implementation. + cudaCheck(cudaGetLastError()); + auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + cudaGetLastError(); + if (status == cudaSuccess) { + layernorm_forward_kernel6<<>>(out, mean, rstd, inp, weight, bias, N, C); + } else { + // fall back to the version without shared memory + const int grid_size_fb = CEIL_DIV(N * WARP_SIZE, block_size); + layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); + } cudaCheck(cudaGetLastError()); } diff --git a/llmc/rand.h b/llmc/rand.h index 99b10fe..63df35f 100644 --- a/llmc/rand.h +++ b/llmc/rand.h @@ -220,13 +220,13 @@ void normal_(float* data, unsigned int numel, float mean, float std, mt19937_sta } } -void random_permutation_with_init(int* data, int numel, mt19937_state* state, int should_init) { - if (should_init) { - for (int i = 0; i < numel; i++) { - data[i] = i; - } +void init_identity_permutation(int *data, int numel) { + for (int i = 0; i < numel; i++) { + data[i] = i; } +} +void random_permutation(int* data, int numel, mt19937_state* state) { for (int i = numel - 1; i > 0; i--) { // pick an index j in [0, i] with equal probability int j = randint32(state) % (i + 1);