diff --git a/dev/cuda/gelu_forward.cu b/dev/cuda/gelu_forward.cu index e69abdc..c9fba9e 100644 --- a/dev/cuda/gelu_forward.cu +++ b/dev/cuda/gelu_forward.cu @@ -21,6 +21,22 @@ version 2 uses the Packed128 data structure #include #include "common.h" +// turn on bf16 as default, done up here for now +//#define ENABLE_BF16 + +#if defined(ENABLE_BF16) +typedef __nv_bfloat16 floatX; +typedef __nv_bfloat16 floatN; +#elif defined(ENABLE_FP16) +typedef half floatX; +typedef half floatN; +#else +typedef float floatX; +typedef float floatN; +#endif + +typedef Packed128 x128; + // ---------------------------------------------------------------------------- // CPU code reference @@ -48,19 +64,19 @@ __global__ void gelu_kernel(float* out, const float* inp, int N) { } // elementwise ops are nice and ez -__global__ void gelu_kernel2(float* out, const float* inp, int N) { - int i = (blockIdx.x * blockDim.x + threadIdx.x) * f128::size; +__global__ void gelu_kernel2(floatX* out, const floatX* inp, int N) { + int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; if (i < N) { - f128 packet_out; - f128 packet_in = load128cs(inp + i); // load and do not keep in cache - for(int k = 0; k < packet_in.size; ++k) { - float xi = packet_in[k]; + x128 packed_out; + x128 packed_inp = load128cs(inp + i); // load and do not keep in cache + for(int k = 0; k < packed_inp.size; ++k) { + float xi = (float)packed_inp[k]; float cube = 0.044715f * xi * xi * xi; - packet_out[k] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))); + packed_out[k] = (float)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)))); } // store instead of storecs (without cache streaming) in case it is useful for the // data to be in the cache for the next operation after this GeLU - store128(out + i, packet_out); + store128(out + i, packed_out); } } @@ -73,25 +89,29 @@ void gelu_forward1(float* out, const float* inp, int N, const int block_size) { cudaCheck(cudaGetLastError()); } -void gelu_forward2(float* out, const float* inp, int N, const int block_size) { - const int grid_size = ceil_div(N, block_size) / 4; +void gelu_forward2(floatX* out, const floatX* inp, int N, const int block_size) { + const int grid_size = ceil_div(N, block_size)/x128::size; gelu_kernel2<<>>(out, inp, N); cudaCheck(cudaGetLastError()); } // kernel version dispatch void gelu_forward(int kernel_num, - float* out, - const float* inp, + floatX* out, + const floatX* inp, int B, int T, int C, int block_size) { switch (kernel_num) { +#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) case 1: gelu_forward1(out, inp, B * T * C, block_size); break; +#endif +#if defined(ENABLE_BF16) case 2: gelu_forward2(out, inp, B * T * C, block_size); break; +#endif default: printf("Invalid kernel number\n"); exit(1); @@ -114,13 +134,6 @@ int main(int argc, char **argv) { float* out = (float*)malloc(B * T * C * sizeof(float)); float* inp = make_random_float(B * T * C); - // move to GPU - float* d_out; - float* d_inp; - cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float))); - cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(float))); - cudaCheck(cudaMemcpy(d_inp, inp, B * T * C * sizeof(float), cudaMemcpyHostToDevice)); - // read kernel_num from command line int kernel_num = 1; if (argc > 1) { @@ -131,6 +144,19 @@ int main(int argc, char **argv) { // first check the correctness of the kernel gelu_forward_cpu(out, inp, B * T * C); + // move to GPU + floatX* d_out; + floatX* d_inp; + cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(floatX))); + cudaCheck(cudaMalloc(&d_inp, B * T * C * sizeof(floatX))); + + floatX* inpX = (floatX*)malloc(B * T * C * sizeof(floatX)); + + for (int i = 0; i < B * T * C; i++) { + inpX[i] = (floatX)inp[i]; + } + + cudaCheck(cudaMemcpy(d_inp, inpX, B * T * C * sizeof(floatX), cudaMemcpyHostToDevice)); // time the kernel at different block sizes int block_sizes[] = {32, 64, 128, 256, 512, 1024}; @@ -138,7 +164,12 @@ int main(int argc, char **argv) { int block_size = block_sizes[j]; printf("Checking block size %d.\n", block_size); gelu_forward(kernel_num, d_out, d_inp, B, T, C, block_size); +#if !defined(ENABLE_BF16) && !defined(ENABLE_FP16) validate_result(d_out, out, "out", B * T * C, 1e-5f); +#endif +#if defined(ENABLE_BF16) +#endif + validate_result(d_out, out, "out", B * T * C, 1e-2f); } printf("All results match. Starting benchmarks.\n\n"); @@ -164,8 +195,9 @@ int main(int argc, char **argv) { // free memory free(out); free(inp); + free(inpX); + cudaCheck(cudaFree(d_out)); cudaCheck(cudaFree(d_inp)); - return 0; } \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index 3913bc5..ce3a8f5 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -227,6 +227,7 @@ struct alignas(16) Packed128 { // short-form typedef typedef Packed128 f128; +typedef Packed128 x128; // load a Packed128 from an aligned memory address template @@ -566,7 +567,8 @@ __global__ void permute_kernel_backward(floatX* dinp, __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; + + 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); @@ -668,11 +670,18 @@ __global__ void residual_forward_kernel(floatX* out, floatX* inp1, floatX* inp2, #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; + int i = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size; 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)))); + x128 packed_out; + x128 packed_inp = load128cs(inp + i); // load and do not keep in cache + for(int k = 0; k < packed_inp.size; ++k) { + float xi =(float)packed_inp[k]; + float cube = 0.044715f * xi * xi * xi; + packed_out[k] = (floatX)(0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube)))); + } + // store instead of storecs (without cache streaming) in case it is useful for the + // data to be in the cache for the next operation after this GeLU + store128(out + i, packed_out); } } @@ -1192,7 +1201,7 @@ void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) { void gelu_forward(floatX* out, const floatX* inp, int N) { const int block_size = 128; - const int grid_size = CEIL_DIV(N, block_size); + const int grid_size = CEIL_DIV(N/x128::size, block_size); gelu_forward_kernel<<>>(out, inp, N); cudaCheck(cudaGetLastError()); }