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

super optimised matmul_bias + cuda streams for better parallelism + fix fused_classifier cache hint + remove remaining cooperative groups

Browse source
This commit is contained in:
ademeure 2024-05-03 16:38:46 +01:00
parent f456a75e3f
commit b58181079a
3 changed files with 217 additions and 104 deletions
Download patch file Download diff file Expand all files Collapse all files

15
profile_gpt2.cu
View file

@ -34,11 +34,22 @@ int main() {
cudaCheck(cudaSetDevice(deviceIdx));
cudaDeviceProp deviceProp;
cudaGetDeviceProperties(&deviceProp, deviceIdx);
cuda_num_SMs = deviceProp.multiProcessorCount;
cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor;
printf("[System]\n");
printf("Device %d: %s\n", deviceIdx, deviceProp.name);
cuda_num_SMs = deviceProp.multiProcessorCount;
cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor;
cuda_arch_major = deviceProp.major;
cuda_arch_minor = deviceProp.minor;
cudaCheck(cudaStreamCreate(&main_stream));
cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming);
cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming);
for (int i = 0; i < num_parallel_streams; i++) {
cudaCheck(cudaStreamCreate(&parallel_streams[i]));
cudaEventCreateWithFlags(&parallel_events[i], cudaEventDisableTiming);
}
// setup cuBLAS and cuBLASLt
cublasCheck(cublasCreate(&cublas_handle));
cublasCheck(cublasLtCreate(&cublaslt_handle));

16
test_gpt2.cu
View file

@ -89,12 +89,22 @@ int main(int argc, char *argv[]) {
cudaCheck(cudaSetDevice(deviceIdx));
cudaDeviceProp deviceProp;
cudaGetDeviceProperties(&deviceProp, deviceIdx);
cuda_num_SMs = deviceProp.multiProcessorCount;
cuda_arch_major = deviceProp.major;
cuda_arch_minor = deviceProp.minor;
printf("[System]\n");
printf("Device %d: %s\n", deviceIdx, deviceProp.name);
cuda_num_SMs = deviceProp.multiProcessorCount;
cuda_threads_per_SM = deviceProp.maxThreadsPerMultiProcessor;
cuda_arch_major = deviceProp.major;
cuda_arch_minor = deviceProp.minor;
cudaCheck(cudaStreamCreate(&main_stream));
cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming);
cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming);
for (int i = 0; i < num_parallel_streams; i++) {
cudaCheck(cudaStreamCreate(&parallel_streams[i]));
cudaEventCreateWithFlags(&parallel_events[i], cudaEventDisableTiming);
}
// setup cuBLAS and cuBLASLt
cublasCheck(cublasCreate(&cublas_handle));
cublasCheck(cublasLtCreate(&cublaslt_handle));

290
train_gpt2.cu
View file

@ -51,8 +51,6 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200),
#include <cuda_runtime.h>
#include <cublasLt.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <nvtx3/nvToolsExt.h>
// Multi-GPU related
@ -156,11 +154,13 @@ int cuda_arch_minor = 0;
int cuda_num_SMs = 0; // for persistent threads where we want 1 threadblock per SM
int cuda_threads_per_SM = 0;
constexpr int num_parallel_streams = 4; // + 1 primary "stream" (+ default stream)
// CUDA streams & events (note: non-timing events, use separate event for timing/profiling!)
constexpr int num_parallel_streams = 2; // + 1 primary "main_stream" (+ default stream)
cudaStream_t parallel_streams[num_parallel_streams];
cudaStream_t stream;
namespace cg = cooperative_groups;
cudaEvent_t parallel_events[num_parallel_streams];
cudaStream_t main_stream;
cudaEvent_t main_event;
cudaEvent_t loss_event; // to make sure fused_classifier has written the losses to the CPU buffer
// convenience macro for calculating grid/block dimensions for kernels
#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
@ -261,7 +261,9 @@ __device__ floatX warpReduceSum(floatX val) {
#endif
using reduction_func_t = float (*) (float);
__device__ float blockReduce(float val, reduction_func_t warp_reduction) {
__device__ float blockReduce(float val, reduction_func_t warp_reduction, bool final_sync=false) {
// two reductions of up to 1024 threads:
// 1) inside warp (shuffle), 2) cross-warp (shared memory), 3) inside warp (shuffle)
__shared__ float shared_val[32];
int lane_id = threadIdx.x % 32;
int warp_id = threadIdx.x / 32;
@ -273,6 +275,10 @@ __device__ float blockReduce(float val, reduction_func_t warp_reduction) {
// same strategy, now reduce across warps
warp_val = (lane_id < num_warps) ? shared_val[lane_id] : 0.0f;
float block_val = warp_reduction(warp_val);
if (final_sync) {
__syncthreads(); // only needed in loops when effectively reusing shared memory etc.
}
return block_val;
}
@ -828,9 +834,12 @@ __global__ void encoder_backward_kernel(floatX* dwte, floatX* 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) {
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
const int warp_size = 32;
int lane_id = threadIdx.x % warp_size;
int warp_id = threadIdx.x / warp_size;
int num_warps = blockDim.x / warp_size;
int idx = blockIdx.x * num_warps + warp_id;
if(idx >= N) { return; } // guard
// the row of input that this group of threads is responsible for
@ -838,30 +847,30 @@ __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __re
// mean
float sum = 0.0f;
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
for (int i = lane_id; i < C; i += warp_size) {
sum += (float)x[i];
}
sum = cg::reduce(warp, sum, cg::plus<float>{});
sum = warpReduceSum(sum);
float m = sum / C;
if(warp.thread_rank() == 0 && mean != nullptr) {
if(lane_id == 0 && mean != nullptr) {
__stcs(mean + idx, (floatX)m);
}
// rstd
sum = 0.0f;
for (int i = warp.thread_rank(); i < C; i += warp.size()) {
for (int i = lane_id; i < C; i += warp_size) {
float diff = (float)x[i] - m;
sum += diff * diff;
}
sum = cg::reduce(warp, sum, cg::plus<float>{});
sum = warpReduceSum(sum);
float s = rsqrtf(sum / C + 1e-5f);
if(warp.thread_rank() == 0 && rstd != nullptr) {
if(lane_id == 0 && rstd != nullptr) {
__stcs(rstd + idx, (floatX)s);
}
// final normalization and scaling by weight/bias
floatX* o = out + idx * C;
for (int c = warp.thread_rank(); c < C; c += warp.size()) {
for (int c = lane_id; c < C; c += warp_size) {
// load and store using the .cs "streaming" hint to the compiler,
// indicating that this data will not be reused soon, and can be streamed through the caches
// this allows the threads to get more cache-hits for the (shared) weight and bias parameters
@ -948,14 +957,17 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons
// directly autoregressive, so we only compute the lower triangular part
// uses the online softmax algorithm
assert(T % 4 == 0);
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
const int warp_size = 32;
int lane_id = threadIdx.x % warp_size;
int warp_id = threadIdx.x / warp_size;
int num_warps = blockDim.x / warp_size;
// micro-optimization: we iterate backwards so that
// after the softmax backward operation completes, the cache retains the
// part of the matrix close to the upper left corner, which benefits the
// matmul operation that immediately follows.
// int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); // forward order
int idx = (gridDim.x - blockIdx.x - 1) * warp.meta_group_size() + warp.meta_group_rank(); // backward order
int idx = (gridDim.x - blockIdx.x - 1) * num_warps + warp_id; // backward order
if(idx >= N * T) {
return;
}
@ -970,7 +982,7 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons
float sumval = 0.0f;
const floatX* x_aligned = reinterpret_cast<const floatX*>(__builtin_assume_aligned(x, 16));
for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) {
for (int i = lane_id; i < pos_by_4; i += warp_size) {
float regarray[4];
#pragma unroll
for (int k = 0; k < 4; ++k) {
@ -986,21 +998,21 @@ __global__ void softmax_forward_kernel5(floatX* out, float inv_temperature, cons
}
}
if(4*pos_by_4 + warp.thread_rank() <= own_pos) {
if(4*pos_by_4 + lane_id <= own_pos) {
float old_maxval = maxval;
maxval = fmaxf(maxval, (float)x[4*pos_by_4 + warp.thread_rank()]);
maxval = fmaxf(maxval, (float)x[4*pos_by_4 + lane_id]);
sumval *= expf(inv_temperature * (old_maxval - maxval));
sumval += expf(inv_temperature * ((float)x[4*pos_by_4 + warp.thread_rank()] - maxval));
sumval += expf(inv_temperature * ((float)x[4*pos_by_4 + lane_id] - maxval));
}
float global_maxval = cg::reduce(warp, maxval, cg::greater<float>{});
float global_maxval = warpReduceMax(maxval);
sumval *= expf(inv_temperature * (maxval - global_maxval));
float sum = cg::reduce(warp, sumval, cg::plus<float>{});
float sum = warpReduceSum(sumval);
float norm = 1.f / sum;
// divide the whole row by the sum
for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) {
for (int i = lane_id; i <= own_pos; i += warp_size) {
// recalculation is faster than doing the round-trip through memory.
float ev = expf(inv_temperature * ((float)__ldcs(x + i) - global_maxval));
__stcs(out + idx * T + i, (floatX)(ev * norm));
@ -1059,19 +1071,53 @@ __global__ void gelu_backward_kernel(floatX* dinp, const floatX* inp, const floa
}
}
__global__ void matmul_backward_bias_kernel5(float* dbias, const floatX* dout, int B, int T, int OC) {
__global__ void matmul_backward_bias_kernel6(float* dbias, const floatX* dout, int B, int T, int OC) {
// note: this kernel reads in floatX, but it writes to float!
// this is because we're using atomics, which are super slow in < fp32 precision on < H100 GPUs
// so the trick is do fp32 atomics to a buffer, and then copy_and_cast the result to floatX
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 += (float)dout[i * OC + oc];
// 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
// blockDim.x is 32 --> single warp being responsible for those 256 OCs
// blockDim.y is 16 --> 16 parallel independent warps processing the same OCs for different BTs
// gridDim.x is OC / 256 --> each block processes 256 OCs
// grimDim.y is max(1, (cuda_num_SMs * cuda_threads_per_SM) / (512 * gridDim.x)); --> fill up the entire GPU!
const int block_size = 512;
const int block_size_x = 32;
const int block_size_y = block_size / block_size_x; // 16
const int OC_per_warp = block_size_x * x128::size; // 256 at BF16
int local_oc = threadIdx.x * x128::size;
int global_oc = blockIdx.x * OC_per_warp + local_oc;
float accumulators[x128::size];
__shared__ float shared[OC_per_warp];
for (int k = 0; k < x128::size; k++) {
accumulators[k] = 0.0f;
}
int thread_id = threadIdx.y * block_size_x + threadIdx.x;
for (int i = thread_id; i < OC_per_warp; i += block_size) {
shared[i] = 0.0f;
}
__syncthreads();
for (int i = blockIdx.y*block_size_y + threadIdx.y; i < B * T; i += gridDim.y*block_size_y) {
x128 packed_dout = load128(dout + global_oc + i*OC);
for (int k = 0; k < x128::size; k++) {
//printf("%d: %f + %f\n", oc, accumulators[k], (float)packed_dout[k]);
accumulators[k] += (float)packed_dout[k];
}
//__syncthreads(); // keep block synchronised to maximise memory locality (?)
}
for (int k = 0; k < x128::size; k++) {
atomicAdd(shared + local_oc + k, accumulators[k]);
}
__syncthreads();
if (threadIdx.y == 0) {
for (int i = threadIdx.x; i < OC_per_warp; i += block_size_x) {
//printf("%d => %f\n", i, shared[i]);
atomicAdd(dbias + i + blockIdx.x*OC_per_warp, shared[i]);
}
}
// and atomically add everything together. atomics within one block are conflict-free!
atomicAdd(dbias + oc, sum);
}
__global__ void layernorm_backward_kernel7(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch,
@ -1169,9 +1215,6 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl
int B, int T, int C, float scale) {
constexpr const int BlockSize = 256;
constexpr int T_per_block = 4;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
__shared__ float block_acc[32];
int idx = blockIdx.y;
// go through blocks in reverse order, so the slowest block starts first
@ -1181,10 +1224,6 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl
datt += idx * T * T;
dpreatt += idx * T * T;
if (warp.meta_group_rank() == 0) {
block_acc[warp.thread_rank()] = 0;
}
for(int to = 0; to < T_per_block; ++to) {
int t = t0 - to;
if(t < 0) return;
@ -1193,15 +1232,13 @@ __global__ void softmax_autoregressive_backward_kernel(floatX* dpreatt, const fl
floatX* dpreatt_bth = dpreatt + t * T;
float local_sum = 0;
for (int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) {
for (int t2 = threadIdx.x; t2 <= t; t2 += BlockSize) {
local_sum += (float)att_bth[t2] * (float)datt_bth[t2];
}
block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus<float>{});
block.sync();
local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
local_sum = blockReduce(local_sum, warpReduceSum);
for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) {
for (int t3 = threadIdx.x; t3 <= t; t3 += BlockSize) {
// don't touch the cache. Some parts will still be here from the previous loop, and
// we want to exploit those.
float acc = (float)__ldcs(att_bth + t3) * ((float)__ldcs(datt_bth + t3) - local_sum);
@ -1264,7 +1301,7 @@ __device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, i
// 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+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
x128 packed_x = load128(x + i * x128::size); // try to keep in cache until next read
for(int k = 0; k < packed_x.size; ++k) {
if (i*x128::size+k >= V) { // bounds checking against real V
continue;
@ -1276,14 +1313,6 @@ __device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, i
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;
// Block Max Reduction -> Maths -> Block Sum Reduction
float block_maxval = blockReduce(thread_maxval, warpReduceMax);
@ -1299,7 +1328,7 @@ __device__ SoftmaxParams prepare_softmax_blockwide(int idx, const floatX* inp, i
__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) {
int idx = blockIdx.x;
int idx = gridDim.x - (blockIdx.x+1); // reverse order for cache hits on matmul data
int ix = targets[idx];
// softmax (reading B * T * V, same logits read again below, hopefully still in cache)
@ -1364,7 +1393,7 @@ void encoder_forward(floatX* out,
const int block_size = 256;
const int N = B * T * C;
const int grid_size = CEIL_DIV(N, (int)(block_size * x128::size));
encoder_forward_kernel3<<<grid_size, block_size, 0, stream>>>(out, inp, wte, wpe, B, T, C);
encoder_forward_kernel3<<<grid_size, block_size, 0, main_stream>>>(out, inp, wte, wpe, B, T, C);
cudaCheck(cudaGetLastError());
}
@ -1375,7 +1404,7 @@ void encoder_backward(floatX* dwte, floatX* dwpe,
const int N = B * T * C;
const int block_size = 256;
const int grid_size = CEIL_DIV(N, block_size);
encoder_backward_kernel<<<grid_size, block_size, 0, stream>>>(dwte, dwpe, dout, inp, B, T, C);
encoder_backward_kernel<<<grid_size, block_size, 0, main_stream>>>(dwte, dwpe, dout, inp, B, T, C);
cudaCheck(cudaGetLastError());
}
@ -1386,7 +1415,7 @@ void layernorm_forward(floatX* out, floatX* mean, floatX* rstd,
const int block_size = 512;
const int N = B * T;
const int grid_size = CEIL_DIV(N * 32, block_size);
layernorm_forward_kernel3<<<grid_size, block_size, 0, stream>>>(out, mean, rstd, inp, weight, bias, N, C);
layernorm_forward_kernel3<<<grid_size, block_size, 0, main_stream>>>(out, mean, rstd, inp, weight, bias, N, C);
cudaCheck(cudaGetLastError());
}
@ -1494,7 +1523,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
v = qkvr + 2 * B * T * C;
int total_threads = B * NH * T * HS;
int num_blocks = CEIL_DIV(total_threads, block_size);
permute_kernel<<<num_blocks, block_size, 0, stream>>>(q, k, v, inp, B, T, NH, HS);
permute_kernel<<<num_blocks, block_size, 0, main_stream>>>(q, k, v, inp, B, T, NH, HS);
cudaCheck(cudaGetLastError());
// IMPORTANT: alpha/beta are FP32 for CUBLAS_COMPUTE_32F even if FP16 inputs/outputs
@ -1522,7 +1551,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
// multiply all elements of preatt elementwise by scale
float scale = 1.0 / sqrtf(HS);
int grid_size = CEIL_DIV(B * NH * T * 32, softmax_block_size);
softmax_forward_kernel5<<<grid_size, softmax_block_size, 0, stream>>>(att, scale, preatt, B * NH, T);
softmax_forward_kernel5<<<grid_size, softmax_block_size, 0, main_stream>>>(att, scale, preatt, B * NH, T);
cudaCheck(cudaGetLastError());
// new approach: first cuBLAS another batched matmul
@ -1543,7 +1572,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
// now unpermute
// y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
num_blocks = CEIL_DIV(B * T * C, block_size);
unpermute_kernel<<<num_blocks, block_size, 0, stream>>>(vaccum, out, B, T, NH, HS);
unpermute_kernel<<<num_blocks, block_size, 0, main_stream>>>(vaccum, out, B, T, NH, HS);
cudaCheck(cudaGetLastError());
}
@ -1551,7 +1580,7 @@ void residual_forward(floatX* out, floatX* inp1, floatX* inp2, int N) {
NVTX_RANGE_FN();
const int block_size = 256;
const int grid_size = CEIL_DIV(N, block_size * x128::size);
residual_forward_kernel<<<grid_size, block_size, 0, stream>>>(out, inp1, inp2, N);
residual_forward_kernel<<<grid_size, block_size, 0, main_stream>>>(out, inp1, inp2, N);
cudaCheck(cudaGetLastError());
}
@ -1559,7 +1588,7 @@ void gelu_forward(floatX* out, const floatX* inp, int N) {
NVTX_RANGE_FN();
const int block_size = 512;
const int grid_size = CEIL_DIV(N, block_size * x128::size);
gelu_forward_kernel2<<<grid_size, block_size, 0, stream>>>(out, inp, N);
gelu_forward_kernel2<<<grid_size, block_size, 0, main_stream>>>(out, inp, N);
cudaCheck(cudaGetLastError());
}
@ -1567,7 +1596,7 @@ void gelu_backward(floatX* dinp, const floatX* inp, const floatX* dout, const in
NVTX_RANGE_FN();
const int block_size = 128;
const int grid_size = CEIL_DIV(N, block_size * x128::size);
gelu_backward_kernel<<<grid_size, block_size, 0, stream>>>(dinp, inp, dout, N);
gelu_backward_kernel<<<grid_size, block_size, 0, main_stream>>>(dinp, inp, dout, N);
cudaCheck(cudaGetLastError());
}
@ -1578,6 +1607,42 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
NVTX_RANGE_FN();
float one = 1.0f;
float zero = 0.0f;
// backward to bias, if given, does a +=
if (dbias != NULL) {
// 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
// blockDim.x is 32 --> single warp being responsible for those 256 OCs
// blockDim.y is 16 --> 16 parallel independent warps processing the same OCs for different BTs
// gridDim.x is OC / 256 --> each block processes 256 OCs
// grimDim.y is max(1, (cuda_num_SMs * cuda_threads_per_SM) / (512 * gridDim.x)); --> fill up the entire GPU!
const int warp_size = 32;
const int block_size = 512;
const int OC_per_warp = warp_size * x128::size; // 256 at BF16
const int block_size_x = 32;
const int block_size_y = block_size / block_size_x; // 16
const int grid_size_x = OC / OC_per_warp; // e.g. 3 horizontal blocks for 768 OCs at BF16
const int grid_size_y = max(1, cuda_threads_per_SM * cuda_num_SMs / (block_size * grid_size_x)); // full GPU!
assert((OC % OC_per_warp) == 0); // there is no bounds checking in the kernel to maximise performance
// Run the memset on a parallel stream then immediately wait on it (runs in parallel with previous kernel)
cudaMemsetAsync(dbias_buffer, 0, OC * sizeof(float), parallel_streams[0]);
cudaEventRecord(parallel_events[0], parallel_streams[0]);
cudaStreamWaitEvent(main_stream, parallel_events[0], 0);
matmul_backward_bias_kernel6<<<dim3(grid_size_x, grid_size_y),
dim3(block_size_x, block_size_y),
OC_per_warp * sizeof(float), main_stream>>>(dbias_buffer, dout, B, T, OC);
cudaEventRecord(main_event, main_stream);
cudaStreamWaitEvent(parallel_streams[0], main_event, 0); // cast_and_add_kernel is dependent on this
// run cast_and_add kernel on separate stream so it can run in parallel with the matmuls
cast_and_add_kernel<<<CEIL_DIV(OC, 256), 256, 0, parallel_streams[0]>>>(dbias, dbias_buffer, OC);
cudaEventRecord(parallel_events[0], parallel_streams[0]);
cudaCheck(cudaGetLastError());
}
// backward to input, uses = in the backward pass (set the gradient)
cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one,
weight, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &zero,
@ -1586,16 +1651,10 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one,
inp, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &one,
dweight, CUBLAS_LOWP, C, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
// backward to bias, if given, does a +=
if (dbias != NULL) {
const int block_size = 128;
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);
cudaMemset(dbias_buffer, 0, OC * sizeof(float));
matmul_backward_bias_kernel5<<<dim3(grid_size_x, grid_size_y), dim3(block_size), 0, stream>>>(dbias_buffer, dout, B, T, OC);
cudaCheck(cudaGetLastError());
cast_and_add_kernel<<<CEIL_DIV(OC, 256), 256, 0, stream>>>(dbias, dbias_buffer, OC);
cudaCheck(cudaGetLastError());
// make sure cast_and_add_kernel is done (hopefully finished in parallel with the matmuls)
cudaStreamWaitEvent(main_stream, parallel_events[0], 0);
}
}
@ -1606,9 +1665,17 @@ void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scr
const int block_size = 1024;
const int grid_size = 1 * cuda_num_SMs;
size_t shared_mem_size = (2 * C + 1) * sizeof(float);
cudaMemset(scratch, 0, (2 * C + 1) * sizeof(float)); // todo - memset in parallel with previous kernels using streams
layernorm_backward_kernel7<<<grid_size, block_size, shared_mem_size, stream>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
cudaMemsetAsync(scratch, 0, (2 * C + 1) * sizeof(float), parallel_streams[0]);
cudaEventRecord(parallel_events[0], parallel_streams[0]);
cudaStreamWaitEvent(main_stream, parallel_events[0], 0); // cast_and_add_kernel is dependent on this
layernorm_backward_kernel7<<<grid_size, block_size, shared_mem_size, main_stream>>>(dinp, dweight, dbias, scratch, dout, inp, weight, mean, rstd, B, T, C);
cudaCheck(cudaGetLastError());
// create a dependency between this kernel and future uses of the scratch buffer (e.g. in matmul_backward)
cudaEventRecord(main_event, main_stream);
cudaStreamWaitEvent(parallel_streams[0], main_event, 0); // cast_and_add_kernel is dependent on this
}
// the sequence of transformations in this compound op is:
@ -1641,7 +1708,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da
// backward through the unpermute operation
int num_blocks = CEIL_DIV(B * T * C, block_size);
unpermute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(scratch, dout, B, T, NH, HS);
unpermute_kernel_backward<<<num_blocks, block_size, 0, main_stream>>>(scratch, dout, B, T, NH, HS);
cudaCheck(cudaGetLastError());
// backward into datt
@ -1657,7 +1724,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da
// backward into preatt
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
softmax_autoregressive_backward_kernel<<<dim3(T / 4, B * NH), 256, 256, stream>>>(dpreatt, datt, att, B, T, C, scale);
softmax_autoregressive_backward_kernel<<<dim3(T / 4, B * NH), 256, 256, main_stream>>>(dpreatt, datt, att, B, T, C, scale);
cudaCheck(cudaGetLastError());
// backward into q
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, alpha_ptr,
@ -1669,7 +1736,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* dpreatt, floatX* da
dk, CUBLAS_LOWP, HS, T * HS, B * NH, CUBLAS_LOWP_COMPUTE, CUBLAS_GEMM_DEFAULT));
// backward into inp
num_blocks = CEIL_DIV(B * NH * T * HS, block_size);
permute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(dinp, dq, dk, dv, B, T, NH, HS);
permute_kernel_backward<<<num_blocks, block_size, 0, main_stream>>>(dinp, dq, dk, dv, B, T, NH, HS);
cudaCheck(cudaGetLastError());
}
@ -1682,7 +1749,7 @@ void fused_classifier3(Type* logits, Type* losses,
const int block_size = 1024;
const int N = B * T;
const int grid_size = N;
fused_classifier_kernel3<<<grid_size, block_size, 512, stream>>>(logits, losses, (Type*)NULL, dlosses, targets, B, T, V, P);
fused_classifier_kernel3<<<grid_size, block_size, 512, main_stream>>>(logits, losses, (Type*)NULL, dlosses, targets, B, T, V, P);
cudaCheck(cudaGetLastError());
}
@ -1794,7 +1861,7 @@ typedef struct {
floatX* lnf; // (B, T, C)
floatX* lnf_mean; // (B, T)
floatX* lnf_rstd; // (B, T)
floatX* losses; // (B, T)
floatX* losses; // (B, T) // todo - no longer used as GPU writes directly to cpu_losses
// adding these two compared to the CPU .c code, needed for attention kernel as buffers
floatX* qkvr; // (L, B, T, 3*C)
// in inference mode, this buffer will store the logits
@ -2062,9 +2129,10 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) {
}
// copy inputs/targets to the model
cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice));
cudaCheck(cudaMemcpyAsync(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice, main_stream));
if (targets != NULL) {
cudaCheck(cudaMemcpy(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice));
cudaCheck(cudaMemcpyAsync(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice, parallel_streams[0]));
cudaEventRecord(parallel_events[0], parallel_streams[0]);
}
// forward pass
@ -2140,16 +2208,18 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) {
// also forward the cross-entropy loss function if we have the targets
if (targets != NULL) {
NvtxRange classifier_and_loss_range("classifier_and_loss");
// wait on memcpy of targets (it definitely finished by now, but better safe than sorry)
cudaStreamWaitEvent(main_stream, parallel_events[0], 0);
// fused classifier: does the forward pass and first part of the backward pass
// we're passing dlosses = NULL, which will default them to 1.0f/(B*T), i.e. uniform loss
fused_classifier3(acts.output, acts.losses, (floatX*)NULL, model->targets, B, T, V, Vp);
// for convenience also evaluate the mean loss (TODO re-think this compute+sync point)
// move the (B,T) losses to CPU
cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(floatX), cudaMemcpyDeviceToHost));
float mean_loss = 0.0f;
for (int i=0; i<B*T; i++) { mean_loss += (float)(model->cpu_losses[i]); }
mean_loss /= B*T;
model->mean_loss = mean_loss;
fused_classifier3(acts.output, model->cpu_losses, (floatX*)NULL, model->targets, B, T, V, Vp);
// the GPU now writes the losses directly to the CPU buffer allocated with cudaMallocHost()
// we accumulate cpu_losses at the end of gpt2_backward() waiting on this event
cudaEventRecord(loss_event, main_stream);
// reset mean_loss here so gpt2_backward() knows we have targets
model->mean_loss = 0.0f;
} else {
// if we don't have targets, we don't have loss
model->mean_loss = -1.0f;
@ -2158,8 +2228,12 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) {
void gpt2_zero_grad(GPT2 *model) {
NVTX_RANGE_FN();
if (model->grads_acts_memory != NULL) { cudaCheck(cudaMemset(model->grads_acts_memory, 0, model->num_grad_acts * sizeof(floatX))); }
if (model->grads_memory != NULL) { cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX))); }
if (model->grads_memory != NULL) {
cudaCheck(cudaMemsetAsync(model->grads_memory, 0, model->num_parameters * sizeof(floatX), parallel_streams[0]));
}
// Allow this to run in parallel with forward pass, but create a dependency with everything after (backwards pass)
cudaEventRecord(parallel_events[0], parallel_streams[0]);
cudaStreamWaitEvent(main_stream, parallel_events[0], 0);
}
void gpt2_backward(GPT2 *model) {
@ -2204,6 +2278,15 @@ void gpt2_backward(GPT2 *model) {
ActivationTensors acts = model->acts;
GradActTensors grads_acts = model->grads_acts;
// reset residual stream gradients (put here to work with gradient accumulation)
cudaCheck(cudaMemsetAsync(model->grads_acts.residual3, 0, B * T * C * sizeof(floatX), parallel_streams[0]));
// allow the memset to run in parallel with the forward pass, but create a dependency with everything after
cudaEventRecord(parallel_events[0], parallel_streams[0]);
cudaStreamWaitEvent(main_stream, parallel_events[0], 0);
// synchronise the other way around as well to avoid read-after-write hazards for the scratch buffer
cudaEventRecord(main_event, main_stream);
cudaStreamWaitEvent(parallel_streams[0], main_event, 0); // cast_and_add_kernel is dependent on this
// re-use the output buffer of the forward pass as a scratchpad during backward pass
float* scratchF = (float*)acts.output;
@ -2292,6 +2375,11 @@ void gpt2_backward(GPT2 *model) {
layernorm_backward(dresidual, dl_ln1w, dl_ln1b, scratchF, dl_btc, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C);
}
encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C);
// accumulate the loss, this was calculated at the end of gpt2_forward()
cudaCheck(cudaEventSynchronize(loss_event)); // hopefully finished long ago
for (int i=0; i<B*T; i++) { model->mean_loss += (float)(model->cpu_losses[i]); }
model->mean_loss /= B*T;
}
// Compute a mean of a single CPU value across all GPU processes. No-op when multi-GPU is disabled.
@ -2318,7 +2406,7 @@ void gpt2_multi_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) {
model->num_parameters,
ncclFloatX, ncclAvg,
multi_gpu_config->nccl_comm,
// use 0 for default stream (all other computations use this stream)
// use 0 for default stream (always implicitly synchronised)
/*stream=*/0));
#endif
}
@ -2338,7 +2426,7 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo
if (model->use_master_weights == 1) {
// allocate one more buffer to keep the master copy of weights as float, and copy the weights over
cudaCheck(cudaMalloc((void**)&model->master_weights, model->num_parameters * sizeof(float)));
copy_and_cast_kernel<<<CEIL_DIV(model->num_parameters, 512), 512, 0, stream>>>(model->master_weights, (floatX*)model->params_memory, model->num_parameters);
copy_and_cast_kernel<<<CEIL_DIV(model->num_parameters, 512), 512, 0, main_stream>>>(model->master_weights, (floatX*)model->params_memory, model->num_parameters);
cudaCheck(cudaGetLastError());
printf0("allocated %zu MiB for master copy of params\n", (model->num_parameters * sizeof(float)) >> 20);
}
@ -2349,7 +2437,7 @@ void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, flo
float beta1_correction = 1.0f - powf(beta1, t);
float beta2_correction = 1.0f - powf(beta2, t);
unsigned int seed = random_u32(&model->rng_state);
adamw_kernel3<<<num_blocks, block_size, 0, stream>>>((floatX*)model->params_memory, model->master_weights,
adamw_kernel3<<<num_blocks, block_size, 0, main_stream>>>((floatX*)model->params_memory, model->master_weights,
(floatX*)model->grads_memory, model->m_memory, model->v_memory,
model->num_parameters,
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay, seed);
@ -2593,9 +2681,12 @@ int main(int argc, char *argv[]) {
cuda_arch_major = deviceProp.major;
cuda_arch_minor = deviceProp.minor;
cudaCheck(cudaStreamCreate(&stream));
cudaCheck(cudaStreamCreate(&main_stream));
cudaEventCreateWithFlags(&main_event, cudaEventDisableTiming);
cudaEventCreateWithFlags(&loss_event, cudaEventDisableTiming);
for (int i = 0; i < num_parallel_streams; i++) {
cudaCheck(cudaStreamCreate(&parallel_streams[i]));
cudaEventCreateWithFlags(&parallel_events[i], cudaEventDisableTiming);
}
// set up cuBLAS and cuBLASLt
@ -2768,6 +2859,7 @@ int main(int argc, char *argv[]) {
}
gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1);
// todo - move or double-buffer all of this timing logic to avoid idling the GPU at this point!
cudaEventRecord(end);
float time_elapsed_ms;
cudaCheck(cudaEventSynchronize(end)); // wait for the end event to finish to get correct timings