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

move encoder, kernel 1/N

Browse source
This commit is contained in:
Andrej Karpathy 2024-06-03 21:21:28 +00:00
parent 3392671826
commit 93546aed01
2 changed files with 133 additions and 129 deletions
Download patch file Download diff file Expand all files Collapse all files

131
llmc/encoder.cuh Normal file
View file

@ -0,0 +1,131 @@
#include "cuda_common.h"
#include "cuda_utils.cuh"
__global__ void encoder_forward_kernel3(floatX* out,
const int* inp, const floatX* wte, const floatX* wpe,
int B, int T, int C) {
int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
int N = B * T * C;
if (idx >= N) { return; }
int bt = idx / C;
int b = bt / T;
int t = bt % T;
int c = idx % C;
int ix = inp[b * T + t];
floatX* out_btc = out + b * T * C + t * C + c;
const floatX* wte_ix = wte + ix * C + c;
const floatX* wpe_tc = wpe + t * C + c;
x128 packed_out;
x128 wte128 = load128cs(wte_ix);
x128 wpe128 = load128cs(wpe_tc);
for (int k = 0; k < x128::size; k++) {
packed_out[k] = (floatX)((float)wte128[k] + (float)wpe128[k]);
}
store128(out_btc, packed_out);
}
template <int BLOCK_SIZE=256>
__global__ void wte_backward_kernel(floatX* dwte,
const int4* bucket_info, const int* workload_indices, const floatX* dout, const int* inp,
unsigned int seed, int B, int T, int C) {
// In order to be deterministic, we preprocess the inputs on the cpu into "buckets"
// Each bucket corresponds to (WARP_SIZE * x128::size) channels for a single vocabulary token
// Each thread handles x128::size channels, e.g. 256 per warp for BF16
// Each block handles (BLOCK_SIZE / WARP_SIZE) elements in a single bucket in parallel
// If a bucket has less than 8 elements, some warps will return immediately
// If a bucket has more than 8 elements, we will loop over all of them
// The buckets are sorted on the CPU so the largest buckets start 1st
int bucket = blockIdx.x;
int warp_id = threadIdx.x / WARP_SIZE;
int lane_id = threadIdx.x % WARP_SIZE;
int c_per_warp = WARP_SIZE * x128::size;
int bucket_start_idx = bucket_info[bucket].x;
int bucket_size = bucket_info[bucket].y;
int bucket_ix = bucket_info[bucket].z;
int c = bucket_info[bucket].w * c_per_warp + (lane_id * x128::size);
// Each thread handles "x128::size" channels, so at fp8, each warp would handle 512 channels
// If C is not a multiple of this (e.g. 768), some buckets/c_groups cannot use the entire warp
if (c >= C) { return; }
// Exit early if this is a small bucket and this warp doesn't have any items to process
if (warp_id >= bucket_size) { return; }
float accum[x128::size] = {0.0f};
__shared__ float accum_shared[x128::size * BLOCK_SIZE];
for(int item = warp_id; item < bucket_size; item += BLOCK_SIZE/WARP_SIZE) {
int bt = workload_indices[bucket_start_idx + item];
const floatX* dout_btc = dout + bt * C + c;
x128 packed_inp1 = load128cs(dout_btc);
for (int k = 0; k < packed_inp1.size; k++) {
accum[k] += (float)packed_inp1[k];
}
}
if (warp_id != 0) {
// we accumulate into warp 0, so only the other warps need to write to shared memory
for (int k = 0; k < x128::size; k++) {
accum_shared[threadIdx.x + k * BLOCK_SIZE] = accum[k];
}
return; // only warp 0 is needed after writing to shared memory
}
// Read dwte for warp 0 even if other warps are not finished yet to maximise latency tolerance
floatX* dwte_ix = dwte + bucket_ix * C + c;
x128 packed_in_out = load128(dwte_ix);
// note: threads which have returned are considered synchronised by CUDA so no risk of deadlock
__syncthreads();
// Accumulate into warp 0's registers by reading the values of the other warps in shared memory
for (int i = threadIdx.x+WARP_SIZE; i < min(BLOCK_SIZE, bucket_size*WARP_SIZE); i += WARP_SIZE) {
for (int k = 0; k < x128::size; k++) {
accum[k] += accum_shared[i + k * BLOCK_SIZE];
}
}
// Add the result to dwte and write back to global memory (read-modify-write)
for (unsigned int k = 0; k < x128::size; k++) {
// We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic
stochastic_rounding(accum[k] + (float)packed_in_out[k], &packed_in_out[k], seed + k);
}
store128(dwte_ix, packed_in_out);
}
__global__ void wpe_backward_kernel(floatX* dwpe,
const floatX* dout, const int* inp,
int B, int T, int C, unsigned int seed) {
// Each thread handles x128::size "channel positions", e.g. 256 per warp for BF16
// For gpt2-124M BF16, C=768 and T=1024, so 3 warps per channel and 3072 warps in total
// For each "channel position" we sum the gradients for every batch at that C/T element
// This way each dwte element is only updated once, and the kernel is fully deterministic!
// The previous kernel was not deterministic, as batches were aggregated with atomicAdd
int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
if (idx >= T * C) { return; }
// if C is not a multiple of WARP_SIZE*x128::size, it's OK for some warps to handle multiple t
int t = idx / C;
int c = idx % C;
float accum[x128::size] = {0.0f};
for (int b = 0; b < B; b++) {
x128 packed_dout = load128cs(dout + (b * T * C) + (t * C) + c); // will never be read again
for (int k = 0; k < x128::size; k++) {
accum[k] += (float)packed_dout[k];
}
}
floatX* dwpe_tc = dwpe + (t * C) + c;
x128 packed_dwpe = load128(dwpe_tc);
for (unsigned int k = 0; k < x128::size; k++) {
// We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic
stochastic_rounding(accum[k] + (float)packed_dwpe[k], &packed_dwpe[k], seed + k);
}
store128(dwpe_tc, packed_dwpe);
}

131
train_gpt2.cu
View file

@ -46,6 +46,8 @@ GPT-2 Transformer Neural Net training loop. See README.md for usage.
// defines: CUBLAS_LOWP, cublasCheck, cublaslt_workspace_size, cublaslt_workspace
// defines: cublas_compute, cublaslt_handle, cublas_handle
#include "llmc/cublas_common.h"
// defines: encoder_forward_kernel3, wte_backward_kernel, wpe_backward_kernel
#include "llmc/encoder.cuh"
// defines: create_cudnn, destroy_cudnn, attention_forward_cudnn, attention_backward_cudnn
#ifdef ENABLE_CUDNN
#include "llmc/cudnn_att.h"
@ -245,135 +247,6 @@ gradients have to add. We make sure that we do a += as necessary.
E.g., the layernorms are connected to the residuals so we += in layernorm backward.
*/
__global__ void encoder_forward_kernel3(floatX* out,
const int* inp, const floatX* wte, const floatX* wpe,
int B, int T, int C) {
int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
int N = B * T * C;
if (idx >= N) { return; }
int bt = idx / C;
int b = bt / T;
int t = bt % T;
int c = idx % C;
int ix = inp[b * T + t];
floatX* out_btc = out + b * T * C + t * C + c;
const floatX* wte_ix = wte + ix * C + c;
const floatX* wpe_tc = wpe + t * C + c;
x128 packed_out;
x128 wte128 = load128cs(wte_ix);
x128 wpe128 = load128cs(wpe_tc);
for (int k = 0; k < x128::size; k++) {
packed_out[k] = (floatX)((float)wte128[k] + (float)wpe128[k]);
}
store128(out_btc, packed_out);
}
template <int BLOCK_SIZE=256>
__global__ void wte_backward_kernel(floatX* dwte,
const int4* bucket_info, const int* workload_indices, const floatX* dout, const int* inp,
unsigned int seed, int B, int T, int C) {
// In order to be deterministic, we preprocess the inputs on the cpu into "buckets"
// Each bucket corresponds to (WARP_SIZE * x128::size) channels for a single vocabulary token
// Each thread handles x128::size channels, e.g. 256 per warp for BF16
// Each block handles (BLOCK_SIZE / WARP_SIZE) elements in a single bucket in parallel
// If a bucket has less than 8 elements, some warps will return immediately
// If a bucket has more than 8 elements, we will loop over all of them
// The buckets are sorted on the CPU so the largest buckets start 1st
int bucket = blockIdx.x;
int warp_id = threadIdx.x / WARP_SIZE;
int lane_id = threadIdx.x % WARP_SIZE;
int c_per_warp = WARP_SIZE * x128::size;
int bucket_start_idx = bucket_info[bucket].x;
int bucket_size = bucket_info[bucket].y;
int bucket_ix = bucket_info[bucket].z;
int c = bucket_info[bucket].w * c_per_warp + (lane_id * x128::size);
// Each thread handles "x128::size" channels, so at fp8, each warp would handle 512 channels
// If C is not a multiple of this (e.g. 768), some buckets/c_groups cannot use the entire warp
if (c >= C) { return; }
// Exit early if this is a small bucket and this warp doesn't have any items to process
if (warp_id >= bucket_size) { return; }
float accum[x128::size] = {0.0f};
__shared__ float accum_shared[x128::size * BLOCK_SIZE];
for(int item = warp_id; item < bucket_size; item += BLOCK_SIZE/WARP_SIZE) {
int bt = workload_indices[bucket_start_idx + item];
const floatX* dout_btc = dout + bt * C + c;
x128 packed_inp1 = load128cs(dout_btc);
for (int k = 0; k < packed_inp1.size; k++) {
accum[k] += (float)packed_inp1[k];
}
}
if (warp_id != 0) {
// we accumulate into warp 0, so only the other warps need to write to shared memory
for (int k = 0; k < x128::size; k++) {
accum_shared[threadIdx.x + k * BLOCK_SIZE] = accum[k];
}
return; // only warp 0 is needed after writing to shared memory
}
// Read dwte for warp 0 even if other warps are not finished yet to maximise latency tolerance
floatX* dwte_ix = dwte + bucket_ix * C + c;
x128 packed_in_out = load128(dwte_ix);
// note: threads which have returned are considered synchronised by CUDA so no risk of deadlock
__syncthreads();
// Accumulate into warp 0's registers by reading the values of the other warps in shared memory
for (int i = threadIdx.x+WARP_SIZE; i < min(BLOCK_SIZE, bucket_size*WARP_SIZE); i += WARP_SIZE) {
for (int k = 0; k < x128::size; k++) {
accum[k] += accum_shared[i + k * BLOCK_SIZE];
}
}
// Add the result to dwte and write back to global memory (read-modify-write)
for (unsigned int k = 0; k < x128::size; k++) {
// We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic
stochastic_rounding(accum[k] + (float)packed_in_out[k], &packed_in_out[k], seed + k);
}
store128(dwte_ix, packed_in_out);
}
__global__ void wpe_backward_kernel(floatX* dwpe,
const floatX* dout, const int* inp,
int B, int T, int C, unsigned int seed) {
// Each thread handles x128::size "channel positions", e.g. 256 per warp for BF16
// For gpt2-124M BF16, C=768 and T=1024, so 3 warps per channel and 3072 warps in total
// For each "channel position" we sum the gradients for every batch at that C/T element
// This way each dwte element is only updated once, and the kernel is fully deterministic!
// The previous kernel was not deterministic, as batches were aggregated with atomicAdd
int idx = (blockIdx.x * blockDim.x + threadIdx.x) * x128::size;
if (idx >= T * C) { return; }
// if C is not a multiple of WARP_SIZE*x128::size, it's OK for some warps to handle multiple t
int t = idx / C;
int c = idx % C;
float accum[x128::size] = {0.0f};
for (int b = 0; b < B; b++) {
x128 packed_dout = load128cs(dout + (b * T * C) + (t * C) + c); // will never be read again
for (int k = 0; k < x128::size; k++) {
accum[k] += (float)packed_dout[k];
}
}
floatX* dwpe_tc = dwpe + (t * C) + c;
x128 packed_dwpe = load128(dwpe_tc);
for (unsigned int k = 0; k < x128::size; k++) {
// We use stochastic rounding to go from FP32 to BF16 but the seed should be deterministic
stochastic_rounding(accum[k] + (float)packed_dwpe[k], &packed_dwpe[k], seed + k);
}
store128(dwpe_tc, packed_dwpe);
}
__global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd,
const floatX* __restrict__ inp, const floatX* __restrict__ weight,
const floatX* __restrict__ bias, int N, int C) {