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

fixes + bounds checking

Browse source
This commit is contained in:
ademeure 2024-04-18 03:41:27 +01:00
parent 4ee98f5135
commit 3ba2cdc723
1 changed files with 49 additions and 34 deletions
Download patch file Download diff file Expand all files Collapse all files

83
dev/cuda/classifier_fused.cu
View file

@ -154,33 +154,31 @@ __device__ float vec_at(const float4& vec, int index) {
return reinterpret_cast<const float*>(&vec)[index];
}
__device__ SoftmaxParams prepare_softmax2(cg::thread_block_tile<32>& warp,
int idx, const float* inp, int V) {
__device__ SoftmaxParams prepare_softmax_blockwide(cg::thread_block_tile<32>& warp,
int idx, const float* inp, int V, int P) {
// one row of inp, i.e. inp[idx, :] of shape (V,)
// float4 to get 128-bit loads and memory level parallelism
// this is only possible if V is a multiple of 4
const float4* x_vec4 = reinterpret_cast<const float4*>(inp + idx * V);
const float4* x_vec4 = reinterpret_cast<const float4*>(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/4) + (threadIdx.x - blockDim.x); i >= 0; i -= blockDim.x) {
float4 v = x_vec4[i];
for (int i = (V+3)/4 + (threadIdx.x - blockDim.x); i >= 0; i -= blockDim.x) {
float4 v4 = x_vec4[i];
#pragma unroll
for(int k = 0; k < 4; ++k) {
float v = (i*4+k < V) ? vec_at(v4, k) : 0.f; // bounds checking against real V
float old_maxval = thread_maxval;
thread_maxval = fmaxf(thread_maxval, vec_at(v, k));
thread_maxval = fmaxf(thread_maxval, vec_at(v4, k));
thread_sumval *= expf((old_maxval - thread_maxval));
thread_sumval += expf(vec_at(v, k) - thread_maxval);
thread_sumval += expf(vec_at(v4, k) - thread_maxval);
}
}
// reduction in 2 stages: 1) inside warp, 2) between warps
// this results in much cleaner code than using a multi-warp cg::reduce
// todo benchmark to make sure it's faster, possibly too many reductions
// we could do the 2nd set of reductions per-block rather than per-warp
// but that would require an extra __syncthreads() unfortunately
// 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;
@ -203,31 +201,50 @@ __device__ SoftmaxParams prepare_softmax2(cg::thread_block_tile<32>& warp,
return SoftmaxParams{1.f / block_sumval, block_maxval};
}
__global__ void fused_classifier_kernel2(float* dlogits, float* losses,
// Fused forward and backward pass for classifier including softmax, and logit gradients
// Writes to both probs (only for debugging) and dlogits (only for training) are optional
// N.B.: We may want to reuse the logits memory for dlogits, so they should *not* be __restrict__!
__global__ void fused_classifier_kernel2(float* dlogits, float* losses, float* probs,
const float* logits, const float* dlosses, const int* targets,
int B, int T, int V) {
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];
float dloss = dlosses[idx];
auto sp = prepare_softmax2(warp, idx, logits, V);
// softmax (reading B * T * V, same logits read again below, hopefully still in cache)
auto sp = prepare_softmax_blockwide(warp, idx, logits, V);
// calculate the probability needed for the loss and update.
// single-threaded
// calculate the probability needed for the loss and update (single-threaded)
if(threadIdx.x == 0) {
float prob = expf(logits[idx * V + ix] - sp.Offset) * sp.Scale;
float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale;
losses[idx] = -logf(prob);
}
// calculate all the gradients
for (int i = threadIdx.x; i < V; i += blockDim.x) {
float prob = expf(__ldcs(&logits[idx * V + i]) - sp.Offset) * sp.Scale;
float indicator = i == ix ? 1.0f : 0.0f;
dlogits[idx * V + i] = (prob - indicator) * dloss;
// calculate the gradients directly, saves bandwidth from probs during training
// but also supports writing probs for inference-only and debugging
float dloss = dlosses ? dlosses[idx] : 0.f;
const float4* logits_vec4 = reinterpret_cast<const float4*>(logits + idx * P);
for (int i = threadIdx.x; i < (V+3)/4; 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
float4 v4 = __ldcs(&logits_vec4[i]);
#pragma unroll
for(int k = 0; k < 4; ++k) {
int element = i*4 + k;
float prob = expf(vec_at(v4, k) - sp.Offset) * sp.Scale;
prob = (element < V) ? prob : 0.f; // bounds checking against real V
// this kernel is DRAM limited so cost of inner branch is ~zero
if (probs) {
probs[idx * P + element] = prob;
}
if (dlogits) {
float indicator = element == ix ? 1.0f : 0.0f;
dlogits[idx * P + element] = (prob - indicator) * dloss;
}
}
}
}
@ -239,16 +256,16 @@ void fused_classifier1(float* dlogits, float* losses,
int B, int T, int V, int P, int block_size) {
const int N = B * T;
const int grid_size = N / (block_size / 32);
fused_classifier_kernel1<<<grid_size, block_size>>>(dlogits, losses, logits, dlosses, targets, B, T, V);
fused_classifier_kernel1<<<grid_size, block_size>>>(dlogits, losses, logits, dlosses, targets, B, T, V, P);
cudaCheck(cudaGetLastError());
}
void fused_classifier2(float* dlogits, float* losses,
const float* logits, const float* dlosses, const int* targets,
int B, int T, int V, int block_size) {
int B, int T, int V, int P, int block_size) {
const int N = B * T;
const int grid_size = N;
fused_classifier_kernel<<<grid_size, block_size>>>(dlogits, losses, logits, dlosses, targets, B, T, V);
fused_classifier_kernel2<<<grid_size, block_size>>>(dlogits, losses, NULL, logits, dlosses, targets, B, T, V, P);
cudaCheck(cudaGetLastError());
}
@ -260,11 +277,7 @@ void fused_classifier(int kernel_num, float* dlogits, float* losses,
fused_classifier1(dlogits, losses, logits, dlosses, targets, B, T, V, P, block_size);
break;
case 2:
if((V % 4) != 0) {
printf("V needs to be a multiple of 4 for this kernel to work!\n");
exit(EXIT_FAILURE);
}
fused_classifier2(dlogits, losses, logits, dlosses, targets, B, T, V, block_size);
fused_classifier2(dlogits, losses, logits, dlosses, targets, B, T, V, P, block_size);
break;
default:
printf("Invalid kernel number\n");
@ -280,6 +293,8 @@ int main(int argc, char **argv) {
int B = 8;
int T = 1024;
int V = 50257;
// padded size
int P = (V + 63) & ~63;
int deviceIdx = 0;
cudaCheck(cudaSetDevice(deviceIdx));