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

Fused classifier: use Packed128 and delete cg dep

Browse source
This commit is contained in:
Andrej Karpathy 2024-05-02 18:08:27 +00:00
commit 2feb9ffcfb
2 changed files with 186 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

130
dev/cuda/classifier_fused.cu
View file

@ -83,6 +83,25 @@ void crossentropy_softmax_backward_cpu(float* dlogits,
}
}
// ----------------------------------------------------
// Kernel Utils
// warp-level reduction for summing values
__device__ float warpReduceSum(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
}
return val;
}
// warp-level reduction for finding the maximum value
__device__ float warpReduceMax(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset));
}
return val;
}
// ----------------------------------------------------------------------------
// GPU kernels
@ -362,6 +381,105 @@ __global__ void fused_classifier_kernel3(float* dlogits, float* losses, float* p
}
}
__device__ SoftmaxParams prepare_softmax_blockwide2(int idx, const float* inp, int V, int P) {
// one row of inp, i.e. inp[idx, :] of shape (V,)
const float* x = inp + idx * P;
float thread_maxval = -INFINITY;
float thread_sumval = 0.0f;
// do the loop in reverse to maximise probability of L2 cache hits
// so even small L2s get some hits on the 2nd read of the same thread
for (int i = (V+3)/4 + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {
f128 packed_x = load128cs(x + i * f128::size); // load and do not keep in cache
for(int k = 0; k < packed_x.size; ++k) {
if (i*4+k >= V) { // bounds checking against real V
continue;
}
float v = (float)packed_x[k];
float old_maxval = thread_maxval;
thread_maxval = fmaxf(thread_maxval, v);
thread_sumval *= expf((old_maxval - thread_maxval));
thread_sumval += expf(v - thread_maxval);
}
}
// two reductions of up to 1024 threads:
// 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)
// this results in much cleaner assembly than a multi-warp cg::reduce
__shared__ float shared_maxval[32];
__shared__ float shared_sumval[32];
int num_warps = blockDim.x / 32;
int warp_id = threadIdx.x / 32;
int lane_id = threadIdx.x % 32;
// reduce maxval within each warp
float warp_maxval = warpReduceMax(thread_maxval);
// thread 0 in each warp writes to shared memory
if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; }
__syncthreads();
// each thread now loads the maxval across previous warps
// if the thread is "out of range" of data, use -FLT_MAX as the maxval
warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX;
// now reduce the maxval among the warp threads
float block_maxval = warpReduceMax(warp_maxval);
// each thread uses maxval to scale sumval to avoid numerical instability / overflow
thread_sumval *= expf(thread_maxval - block_maxval);
// (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory
float warp_sumval = warpReduceSum(thread_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});
if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; }
__syncthreads();
// same strategy, now reduce sumval across warps
warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f;
float block_sumval = warpReduceSum(warp_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});
// return the softmax parameters
return SoftmaxParams{1.f / block_sumval, block_maxval};
}
// same as 2 but not using float4
__global__ void fused_classifier_kernel4(float* dlogits, float* losses, float* probs,
const float* logits, const float* dlosses, const int* targets,
int B, int T, int V, int P) {
int idx = blockIdx.x;
int ix = targets[idx];
// softmax (reading B * T * V, same logits read again below, hopefully still in cache)
SoftmaxParams sp = prepare_softmax_blockwide2(idx, logits, V, P);
// calculate the probability needed for the loss and update (single-threaded)
if(threadIdx.x == 0) {
float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale;
losses[idx] = -logf(prob);
}
// very sensible default for dlosses is 1/(B*T), which is the uniform loss
float dloss = dlosses != NULL ? dlosses[idx] : 1.0f / (B*T);
// calculate the gradients directly, saves bandwidth from probs during training
// but also supports writing probs for inference-only and debugging
const float* logits_vec = logits + idx * P;
for (int i = threadIdx.x; i < (V+f128::size-1)/f128::size; i += blockDim.x) {
// this is the 2nd read of logits after the one in prepare_softmax2
// this data will never be needed again, so we reduce cache persistence
f128 packed_logits_vec = load128cs(logits_vec + i * f128::size); // load and do not keep in cache
f128 packed_probs;
f128 packed_dlogits;
for(int k = 0; k < packed_logits_vec.size; ++k) {
int element = i*packed_logits_vec.size + k;
if (element >= V) { // bounds checking against real V
continue;
}
float v = packed_logits_vec[k];
float prob = expf(v - sp.Offset) * sp.Scale;
packed_probs[k] = prob;
float indicator = (element == ix) ? 1.0f : 0.0f;
packed_dlogits[k] = (prob - indicator) * dloss;
}
store128(dlogits + idx * P + i * packed_logits_vec.size, packed_dlogits);
if (probs != NULL) {
store128(probs + idx * P + i * packed_logits_vec.size, packed_probs);
}
}
}
// ----------------------------------------------------------------------------
// kernel launcher
@ -395,6 +513,15 @@ void fused_classifier3(float* dlogits, float* losses,
cudaCheck(cudaGetLastError());
}
void fused_classifier4(float* dlogits, float* losses,
const float* logits, const float* dlosses, const int* targets,
int B, int T, int V, int P, int block_size) {
const int N = B * T;
const int grid_size = N;
fused_classifier_kernel4<<<grid_size, block_size>>>(dlogits, losses, NULL, logits, dlosses, targets, B, T, V, P);
cudaCheck(cudaGetLastError());
}
void fused_classifier(int kernel_num, float* dlogits, float* losses,
const float* logits, const float* dlosses, const int* targets,
int B, int T, int V, int P, int block_size) {
@ -408,6 +535,9 @@ void fused_classifier(int kernel_num, float* dlogits, float* losses,
case 3:
fused_classifier3(dlogits, losses, logits, dlosses, targets, B, T, V, P, block_size);
break;
case 4:
fused_classifier4(dlogits, losses, logits, dlosses, targets, B, T, V, P, block_size);
break;
default:
printf("Invalid kernel number\n");
exit(1);

83
train_gpt2.cu
View file

@ -213,6 +213,22 @@ __device__ void atomicAddX(float* addr, float val) {
atomicAdd(addr, val);
}
// warp-level reduction for summing values
__device__ float warpReduceSum(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
}
return val;
}
// warp-level reduction for finding the maximum value
__device__ float warpReduceMax(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset));
}
return val;
}
// ----------------------------------------------------------------------------
// Packed128 data structure, which forces the compiler to use 128-bit loads/stores
// in GPUs that support (the LDG.128 and STS.128 instructions)
@ -1195,8 +1211,7 @@ struct SoftmaxParams {
float Offset;
};
__device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_tile<32>& warp,
int idx, const floatX* inp, int V, int P) {
__device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, int V, int P) {
// same but not float4
// one row of inp, i.e. inp[idx, :] of shape (V,)
@ -1205,14 +1220,19 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til
float thread_sumval = 0.0f;
// do the loop in reverse to maximise probability of L2 cache hits
// so even small L2s get some hits on the 2nd read of the same thread
for (int i = V + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {
float v = (float)x[i];
float old_maxval = thread_maxval;
thread_maxval = fmaxf(thread_maxval, v);
thread_sumval *= expf((old_maxval - thread_maxval));
thread_sumval += expf(v - thread_maxval);
for (int i = (V+x128::size-1)/x128::size + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {
x128 packed_x = load128cs(x + i * x128::size); // load and do not keep in cache
for(int k = 0; k < packed_x.size; ++k) {
if (i*x128::size+k >= V) { // bounds checking against real V
continue;
}
float v = (float)packed_x[k];
float old_maxval = thread_maxval;
thread_maxval = fmaxf(thread_maxval, v);
thread_sumval *= expf((old_maxval - thread_maxval));
thread_sumval += expf(v - thread_maxval);
}
}
// two reductions of up to 1024 threads:
// 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)
// this results in much cleaner assembly than a multi-warp cg::reduce
@ -1223,7 +1243,7 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til
int lane_id = threadIdx.x % 32;
// reduce maxval within each warp
float warp_maxval = cg::reduce(warp, thread_maxval, cg::greater<float>{});
float warp_maxval = warpReduceMax(thread_maxval);
// thread 0 in each warp writes to shared memory
if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; }
__syncthreads();
@ -1231,16 +1251,17 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til
// if the thread is "out of range" of data, use -FLT_MAX as the maxval
warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX;
// now reduce the maxval among the warp threads
float block_maxval = cg::reduce(warp, warp_maxval, cg::greater<float>{});
float block_maxval = warpReduceMax(warp_maxval);
// each thread uses maxval to scale sumval to avoid numerical instability / overflow
thread_sumval *= expf(thread_maxval - block_maxval);
// (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory
float warp_sumval = cg::reduce(warp, thread_sumval, cg::plus<float>{});
float warp_sumval = warpReduceSum(thread_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});
if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; }
__syncthreads();
// same strategy, now reduce sumval across warps
warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f;
float block_sumval = cg::reduce(warp, warp_sumval, cg::plus<float>{});
float block_sumval = warpReduceSum(warp_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});
// return the softmax parameters
return SoftmaxParams{1.f / block_sumval, block_maxval};
}
@ -1250,19 +1271,16 @@ __device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_til
__global__ void fused_classifier_kernel3(floatX* logits, floatX* losses, floatX* probs,
const floatX* dlosses, const int* targets,
int B, int T, int V, int P) {
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
int idx = blockIdx.x;
int ix = targets[idx];
// softmax (reading B * T * V, same logits read again below, hopefully still in cache)
SoftmaxParams sp = prepare_softmax_blockwide_nofloat4(warp, idx, logits, V, P);
SoftmaxParams sp = prepare_softmax_blockwide(idx, logits, V, P);
// calculate the probability needed for the loss and update (single-threaded)
if(threadIdx.x == 0) {
float prob = expf((float)logits[idx * P + ix] - sp.Offset) * sp.Scale;
losses[idx] = (floatX)(-logf(prob));
losses[idx] = (floatX)(-logf(prob));
}
// very sensible default for dlosses is 1/(B*T), which is the uniform loss
@ -1270,18 +1288,29 @@ __global__ void fused_classifier_kernel3(floatX* logits, floatX* losses, floatX*
// 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;
// note that we use the padded dimension P to access data, but we only ever
// modify the elements up to V, ignoring the padded dimensions and leaving them at 0
for (int i = threadIdx.x; i < V; i += blockDim.x) {
for (int i = threadIdx.x; i < (V+x128::size-1)/x128::size; i += blockDim.x) {
// this is the 2nd read of logits after the one in prepare_softmax2
// this data will never be needed again, so we reduce cache persistence
float v = (float)__ldcs(&logits_vec[i]);
float prob = expf(v - sp.Offset) * sp.Scale;
if (probs != NULL) {
probs[idx * P + i] = (floatX)prob;
x128 packed_logits_vec = load128cs(logits_vec + i * x128::size); // load and do not keep in cache
x128 packed_probs;
x128 packed_logits;
for(int k = 0; k < packed_logits_vec.size; ++k) {
int element = i*packed_logits_vec.size + k;
if (element >= V) { // bounds checking against real V
continue;
}
float v = (float)packed_logits_vec[k];
float prob = expf(v - sp.Offset) * sp.Scale;
packed_probs[k] = (floatX)prob;
float indicator = (element == ix) ? 1.0f : 0.0f;
packed_logits[k] = (floatX)((prob - indicator) * dloss);
}
if (logits != NULL){
store128(logits + idx * P + i * packed_logits_vec.size, packed_logits);
}
if (probs != NULL) {
store128(probs + idx * P + i * packed_logits_vec.size, packed_probs);
}
float indicator = (i == ix) ? 1.0f : 0.0f;
logits[idx * P + i] = (floatX)((prob - indicator) * dloss);
}
}