aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

speed up the backward bias kernel by 45% and speed up the full running time by 1%

Browse source
This commit is contained in:
Andrej Karpathy 2024-04-20 22:46:19 +00:00
parent 9b722ce5c0
commit 8488669d25
2 changed files with 278 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

255
dev/cuda/matmul_backward_bias.cu Normal file
View file

@ -0,0 +1,255 @@
/*
Kernels for matmul backward pass bias only.
Compile example:
nvcc -O3 matmul_backward_bias.cu -lineinfo -o matmul_backward_bias
./matmul_backward_bias 1
./matmul_backward_bias 2
./matmul_backward_bias 3
ncu:
sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1
*/
#include <stdio.h>
#include <stdlib.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <omp.h>
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include "common.h"
// ----------------------------------------------------------------------------
// CPU code reference
void matmul_backward_bias_cpu(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight,
int B, int T, int C, int OC) {
for (int o = 0; o < OC; o++) {
double sum = 0.0;
for (int b = 0; b < B; b++) {
for (int t = 0; t < T; t++) {
float* dout_bt = dout + b * T * OC + t * OC;
sum += dout_bt[o];
}
}
dbias[o] = sum;
}
}
// ----------------------------------------------------------------------------
// GPU kernels
__global__ void matmul_backward_bias_kernel1(float* dbias, const float* dout, int B, int T, int OC) {
extern __shared__ float shared[];
int o = blockIdx.x; // range [0, OC)
int tid = threadIdx.x; // range [0, block_size)
int block_size = blockDim.x;
const float* x = dout + o;
// thread coarsening
double sum = 0.0;
for (int i = tid; i < B * T; i += block_size) {
sum += x[i * OC];
}
shared[tid] = (float) sum;
__syncthreads();
// reductions
for (int stride = block_size / 2; stride >= 1; stride /= 2) {
__syncthreads();
if (tid < stride) {
shared[tid] += shared[tid + stride];
}
}
// write the final result (at thread 0) to global memory
if (tid == 0) {
dbias[o] += shared[0];
}
}
// cooperative groups solution, one warp per output channel
__global__ void matmul_backward_bias_kernel2(float* dbias, const float* dout, int B, int T, int OC) {
// dout is (B, T, OC), dbias is (OC)
// e.g. if block_size = 128, then we have 4 warps per block, each in charge of one output channel
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
// meta_group_size is the number of warps in a block (e.g. 4), meta_group_rank is the warp index (0,1,2,3)
int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
if(idx >= OC) { return; }
int BT = B * T; // number of elements to reduce in total, per channel
// first, thread coarsening to sum reduce the problem size from B*T to 32
float sum = 0.0f;
for(int i = warp.thread_rank(); i < BT; i += warp.size()) {
sum += dout[i * OC + idx];
}
// now do a warp-level reduce to get the sum across the 32 threads in this warp
sum = cg::reduce(warp, sum, cg::plus<float>{});
// write the result to output (global memory)
if(warp.thread_rank() == 0) {
dbias[idx] += sum;
}
}
__global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, int B, int T, int OC) {
// dout is (B, T, OC), dbias is (OC)
// in this version of the kernel the entire block of block_size is dedicated to one output channel
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
__shared__ float shared_sum[32]; // block_size max is 1024 = 32 * 32 warps
int BT = B * T; // number of elements to reduce in total, per channel
int num_warps = blockDim.x / 32;
int warp_id = threadIdx.x / 32;
int lane_id = threadIdx.x % 32;
int idx = blockIdx.x; // simply one block per row
// round 1: thread coarsening to reduce the problem size from B*T to 32
float thread_sum = 0.0f;
for(int i = threadIdx.x; i < BT; i += blockDim.x) {
thread_sum += dout[i * OC + idx];
}
// now do a warp-level reduce to get the sum across the 32 threads in each warp
float warp_sum = cg::reduce(warp, thread_sum, cg::plus<float>{});
// store the warp sum in shared memory (we could have lane_id == 0 guard but not needed)
shared_sum[warp_id] = warp_sum;
__syncthreads();
// load results from shared memory to threads, pad with zeros for threads that are out of bounds
warp_sum = (lane_id < num_warps) ? shared_sum[lane_id] : 0.0f;
// now reduce the warp-level reductions
float block_sum = cg::reduce(warp, warp_sum, cg::plus<float>{}); // sum(x)
// write the result to output (global memory)
if(threadIdx.x == 0) {
dbias[idx] += block_sum;
}
}
// ----------------------------------------------------------------------------
// kernel launcher
// version1: simple cuBLAS calls
void matmul_backward_bias1(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
int B, int T, int C, int OC, int block_size) {
dim3 block_dim(block_size);
dim3 grid_dim(OC);
size_t shared_mem_size = block_size * sizeof(float);
matmul_backward_bias_kernel1<<<grid_dim, block_dim, shared_mem_size>>>(dbias, dout, B, T, OC);
}
void matmul_backward_bias2(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
int B, int T, int C, int OC, int block_size) {
// block_size 512 seems best
const int grid_size = ceil_div(OC * 32, block_size);
matmul_backward_bias_kernel2<<<grid_size, block_size>>>(dbias, dout, B, T, OC);
}
void matmul_backward_bias3(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
int B, int T, int C, int OC, int block_size) {
// block_size 256 seems best
matmul_backward_bias_kernel3<<<OC, block_size>>>(dbias, dout, B, T, OC);
}
void matmul_backward_bias(int kernel_num,
float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
int B, int T, int C, int OC, int block_size) {
switch (kernel_num) {
case 1:
matmul_backward_bias1(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
break;
case 2:
matmul_backward_bias2(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
break;
case 3:
matmul_backward_bias3(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
break;
default:
printf("Invalid kernel number\n");
exit(1);
}
}
// ----------------------------------------------------------------------------
int main(int argc, char **argv) {
srand(0);
int B = 8;
int T = 1024;
int C = 768;
int OC = 768 * 4; // expansion of 4, e.g. in the MLP
// set up the device
int deviceIdx = 0;
cudaCheck(cudaSetDevice(deviceIdx));
cudaDeviceProp deviceProp;
cudaGetDeviceProperties(&deviceProp, deviceIdx);
printf("Device %d: %s\n", deviceIdx, deviceProp.name);
// read kernel_num from command line
int kernel_num = 1;
if (argc > 1) {
kernel_num = atoi(argv[1]);
}
printf("Using kernel %d\n", kernel_num);
// create host memory of random numbers
float* dbias = make_zeros_float(OC);
float* dout = make_random_float(B * T * OC);
// move to GPU
float* d_dbias;
float* d_dout;
cudaCheck(cudaMalloc(&d_dbias, OC * sizeof(float)));
cudaCheck(cudaMalloc(&d_dout, B * T * OC * sizeof(float)));
cudaCheck(cudaMemcpy(d_dbias, dbias, OC * sizeof(float), cudaMemcpyHostToDevice));
cudaCheck(cudaMemcpy(d_dout, dout, B * T * OC * sizeof(float), cudaMemcpyHostToDevice));
// ncu debugging / profiling, do a single call
// int block_size_debug;
// if (kernel_num == 1) { block_size_debug = 512;
// } else if (kernel_num == 2) { block_size_debug = 512;
// } else { block_size_debug = 256; }
// printf("kernel %d, block_size %d\n", kernel_num, block_size_debug);
// matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, block_size_debug);
// exit(EXIT_SUCCESS);
int block_sizes[] = {32, 64, 128, 256, 512, 1024};
// calculate the CPU reference
matmul_backward_bias_cpu(NULL, NULL, dbias, dout, NULL, NULL, B, T, C, OC);
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
int block_size = block_sizes[j];
// memset the bias to zero
cudaCheck(cudaMemset(d_dbias, 0, OC * sizeof(float)));
// calculate the GPU version
matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, 128);
// compare
printf("Checking correctness...\n");
validate_result(d_dbias, dbias, "dbias", OC, 1e-3f);
printf("All results match for block_size=%d.\n\n", block_size);
}
// now benchmark the kernel
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
int block_size = block_sizes[j];
float *d_dinp, *d_dweight, *d_inp, *d_weight, *d_ones;
int repeat_times = 2000;
float elapsed_time = benchmark_kernel(repeat_times, matmul_backward_bias, kernel_num,
d_dinp, d_dweight, d_dbias, d_dout, d_inp, d_weight, d_ones,
B, T, C, OC, block_size);
printf("block_size %d time %.4f ms\n", block_size, elapsed_time);
}
// cleanups
free(dbias);
free(dout);
cudaCheck(cudaFree(d_dbias));
cudaCheck(cudaFree(d_dout));
return 0;
}

50
train_gpt2.cu
View file

@ -532,29 +532,27 @@ __global__ void softmax_forward_kernel7(float* out, const float* inp, int N, int
}
}
__global__ void matmul_backward_bias_kernel_faster(float* dbias, const float* dout, int B, int T, int OC) {
extern __shared__ float shared[];
int o = blockIdx.x; // range [0, OC)
int tid = threadIdx.x; // range [0, block_size)
int block_size = blockDim.x;
const float* x = dout + o;
// thread coarsening
double sum = 0.0;
for (int i = tid; i < B * T; i += block_size) {
sum += x[i * OC];
// cooperative groups solution, one warp per output channel
__global__ void matmul_backward_bias_kernel2(float* dbias, const float* dout, int B, int T, int OC) {
// dout is (B, T, OC), dbias is (OC)
// e.g. if block_size = 128, then we have 4 warps per block, each in charge of one output channel
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
// meta_group_size is the number of warps in a block (e.g. 4), meta_group_rank is the warp index (0,1,2,3)
int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
if(idx >= OC) { return; }
int BT = B * T; // number of elements to reduce in total, per channel
// first, thread coarsening to sum reduce the problem size from B*T to 32
float sum = 0.0f;
for(int i = warp.thread_rank(); i < BT; i += warp.size()) {
sum += dout[i * OC + idx];
}
shared[tid] = (float) sum;
__syncthreads();
// reductions
for (int stride = block_size / 2; stride >= 1; stride /= 2) {
__syncthreads();
if (tid < stride) {
shared[tid] += shared[tid + stride];
}
}
// write the final result (at thread 0) to global memory
if (tid == 0) {
dbias[o] += shared[0];
// now do a warp-level reduce to get the sum across the 32 threads in this warp
sum = cg::reduce(warp, sum, cg::plus<float>{});
// write the result to output (global memory)
if(warp.thread_rank() == 0) {
dbias[idx] += sum;
}
}
@ -971,11 +969,9 @@ void matmul_backward(float* dinp, float* dweight, float* dbias,
cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one, inp, C, dout, OC, &one, dweight, C));
// backward to bias, if given, does a +=
if (dbias != NULL) {
const int block_size=512;
dim3 block_dim(block_size);
dim3 grid_dim(OC);
size_t shared_mem_size = block_size * sizeof(float);
matmul_backward_bias_kernel_faster<<<grid_dim, block_dim, shared_mem_size>>>(dbias, dout, B, T, OC);
const int block_size = 512;
const int grid_size = CEIL_DIV(OC * 32, block_size);
matmul_backward_bias_kernel2<<<grid_size, block_size>>>(dbias, dout, B, T, OC);
cudaCheck(cudaGetLastError());
}
}