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

enable bf16

Browse source
This commit is contained in:
Erik Schultheis 2024-05-12 19:42:32 +03:00
parent ec92368f8b
commit 2287da0120
1 changed files with 40 additions and 49 deletions
Download patch file Download diff file Expand all files Collapse all files

89
dev/cuda/matmul_backward_bias.cu
View file

@ -23,7 +23,7 @@ sudo ncu --set full --import-source yes -o bias -f ./matmul_backward_bias 1
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
//#define ENABLE_BF16
#define ENABLE_BF16
#include "common.h"
// ----------------------------------------------------------------------------
@ -49,16 +49,16 @@ void matmul_backward_bias_cpu(float* dinp, float* dweight, float* dbias,
float* dbias_buffer;
__global__ void matmul_backward_bias_kernel1(float* dbias, const float* dout, int B, int T, int OC) {
__global__ void matmul_backward_bias_kernel1(floatX* dbias, const floatX* 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;
const floatX* x = dout + o;
// thread coarsening
float sum = 0.0;
for (int i = tid; i < B * T; i += block_size) {
sum += x[i * OC];
sum += (float)x[i * OC];
}
shared[tid] = sum;
__syncthreads();
@ -71,12 +71,12 @@ __global__ void matmul_backward_bias_kernel1(float* dbias, const float* dout, in
}
// write the final result (at thread 0) to global memory
if (tid == 0) {
dbias[o] += shared[0];
dbias[o] = (float)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) {
__global__ void matmul_backward_bias_kernel2(floatX* dbias, const floatX* 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;
@ -89,7 +89,7 @@ __global__ void matmul_backward_bias_kernel2(float* dbias, const float* dout, in
// 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];
sum += (float)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>{});
@ -99,7 +99,7 @@ __global__ void matmul_backward_bias_kernel2(float* dbias, const float* dout, in
}
}
__global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, int B, int T, int OC) {
__global__ void matmul_backward_bias_kernel3(floatX* dbias, const floatX* 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;
@ -114,7 +114,7 @@ __global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, in
// 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];
thread_sum += (float)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>{});
@ -136,7 +136,7 @@ __global__ void matmul_backward_bias_kernel3(float* dbias, const float* dout, in
// 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
__global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, int B, int T, int OC) {
__global__ void matmul_backward_bias_kernel4(floatX* dbias, const floatX* 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)
@ -147,7 +147,7 @@ __global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, in
// 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 float* dout_col = dout + tl + lane_id;
const floatX* 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
@ -156,7 +156,7 @@ __global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, in
// leading to a coalesced memory access pattern
float dout_sum = 0.0f;
for (int row = warp_id; row < B * T; row += vstep) {
dout_sum += dout_col[row * OC];
dout_sum += (float)dout_col[row * OC];
}
smem[lane_id + warp_id * warpSize] = dout_sum;
__syncthreads();
@ -171,13 +171,13 @@ __global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, in
}
}
__global__ void matmul_backward_bias_kernel5(float* dbias, const float* dout, int B, int T, int OC) {
__global__ void matmul_backward_bias_kernel5(floatX* dbias, const floatX* dout, int B, int T, int OC) {
int oc = blockIdx.x * blockDim.x + threadIdx.x;
if(oc >= OC) return;
float sum = 0.0;
// grid-wide loop for maximum parallelism
for (int i = blockIdx.y; i < B * T; i += gridDim.y) {
sum += dout[i * OC + oc];
sum += (float)dout[i * OC + oc];
}
// and atomically add everything together. atomics within one block are conflict-free!
atomicAdd(dbias + oc, sum);
@ -248,8 +248,7 @@ __global__ void matmul_backward_bias_kernel7(float* dbias, const floatX* dout, i
// kernel launcher
// version1: simple cuBLAS calls
void matmul_backward_bias1(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
void matmul_backward_bias1(floatX* dbias, floatX* dout,
int B, int T, int C, int OC, int block_size) {
dim3 block_dim(block_size);
dim3 grid_dim(OC);
@ -257,42 +256,37 @@ void matmul_backward_bias1(float* dinp, float* dweight, float* dbias,
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,
void matmul_backward_bias2(floatX* dbias, floatX* dout,
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,
void matmul_backward_bias3(floatX* dbias, floatX* dout,
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_bias4(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
void matmul_backward_bias4(floatX* dbias, floatX* dout,
int B, int T, int C, int OC, int block_size) {
assert(OC % 32 == 0); // OC must be divisible by 32 for this kernel
const int grid_size = OC / 32;
matmul_backward_bias_kernel4<<<grid_size, block_size, block_size * sizeof(float)>>>(dbias, dout, B, T, OC);
}
void matmul_backward_bias5(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
void matmul_backward_bias5(floatX* dbias, floatX* dout,
int B, int T, int C, int OC, int block_size) {
const int grid_size_x = ceil_div(OC, block_size);
const int grid_size_y = max(1, cuda_threads_per_SM * cuda_num_SMs / block_size);
matmul_backward_bias_kernel5<<<dim3(grid_size_x, grid_size_y), dim3(block_size)>>>(dbias, dout, B, T, OC);
}
void matmul_backward_bias7(float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
void matmul_backward_bias7(floatX* dbias, floatX* dout,
int B, int T, int C, int OC, int block_size) {
if(block_size < 128) {
block_size = 128;
if(block_size < 256) {
block_size = 256;
}
// Each warp is responsible for 32 * "x128::size" = 256 OCs at BF16 (OC must be a multiple of 256!)
// Block size is 512 threads (16 warps) and we reduce those 16 values into 1 at the end
@ -315,28 +309,26 @@ void matmul_backward_bias7(float* dinp, float* dweight, float* dbias,
cast_and_add_kernel<<<ceil_div(OC, 256), 256, 0>>>(dbias, dbias_buffer, OC);
}
void matmul_backward_bias(int kernel_num,
float* dinp, float* dweight, float* dbias,
float* dout, float* inp, float* weight, float* ones,
void matmul_backward_bias(int kernel_num, floatX* dbias, floatX* dout,
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);
matmul_backward_bias1(dbias, dout, 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);
matmul_backward_bias2(dbias, dout, 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);
matmul_backward_bias3(dbias, dout, B, T, C, OC, block_size);
break;
case 4:
matmul_backward_bias4(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
matmul_backward_bias4(dbias, dout, B, T, C, OC, block_size);
break;
case 5:
matmul_backward_bias5(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
matmul_backward_bias5(dbias, dout, B, T, C, OC, block_size);
break;
case 7:
matmul_backward_bias7(dinp, dweight, dbias, dout, inp, weight, ones, B, T, C, OC, block_size);
matmul_backward_bias7(dbias, dout, B, T, C, OC, block_size);
break;
default:
printf("Invalid kernel number\n");
@ -366,13 +358,13 @@ int main(int argc, char **argv) {
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)));
floatX* d_dbias;
floatX* d_dout;
cudaCheck(cudaMalloc(&d_dbias, OC * sizeof(floatX)));
cudaCheck(cudaMalloc(&d_dout, B * T * OC * sizeof(floatX)));
cudaCheck(cudaMalloc(&dbias_buffer, OC * sizeof(float)));
cudaCheck(cudaMemcpy(d_dbias, dbias, OC * sizeof(float), cudaMemcpyHostToDevice));
cudaCheck(cudaMemcpy(d_dout, dout, B * T * OC * sizeof(float), cudaMemcpyHostToDevice));
cudaCheck(memcpy_convert(d_dbias, dbias, OC));
cudaCheck(memcpy_convert(d_dout, dout, B * T * OC));
// ncu debugging / profiling, do a single call
// int block_size_debug;
@ -391,23 +383,22 @@ int main(int argc, char **argv) {
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)));
cudaCheck(cudaMemset(d_dbias, 0, OC * sizeof(floatX)));
// calculate the GPU version
matmul_backward_bias(kernel_num, NULL, NULL, d_dbias, d_dout, NULL, NULL, NULL, B, T, C, OC, block_size);
matmul_backward_bias(kernel_num, d_dbias, d_dout, B, T, C, OC, block_size);
// compare
printf("Checking correctness...\n");
validate_result(d_dbias, dbias, "dbias", OC, 5e-3f);
float tol = std::is_same_v<floatX, float> ? 5e-3f : 1.0f;
validate_result(d_dbias, dbias, "dbias", OC, tol);
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);
d_dbias, d_dout, B, T, C, OC, block_size);
printf("block_size %d time %.4f ms\n", block_size, elapsed_time);
}