diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4f71900..27ebad6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,6 +1,7 @@ name: Build and test on: + create: workflow_dispatch: push: branches: @@ -138,6 +139,33 @@ jobs: call "C:\\Program Files\\Microsoft Visual Studio\\2022\\Enterprise\\VC\\Auxiliary\\Build\\vcvars64.bat" make-4.4.1\dist\make -j WIN_CI_BUILD=1 train_gpt2fp32cu test_gpt2fp32cu test_gpt2cu train_gpt2cu profile_gpt2cu + build-ubuntu20-04: + runs-on: ubuntu-20.04 + container: + image: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04 + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: System Info + run: | + nvcc --version + g++ --version + + - name: Install cudnn frontend + run: | + apt-get update && apt-get install -y git + git clone https://github.com/NVIDIA/cudnn-frontend.git + + - name: Build FP32 checkpoint + run: make train_gpt2fp32cu test_gpt2fp32cu + + - name: Build FP32 precision + run: PRECISION=FP32 make train_gpt2cu test_gpt2cu profile_gpt2cu + + - name: Build with CUDNN + run: PRECISION=BF16 USE_CUDNN=1 make train_gpt2cu test_gpt2cu profile_gpt2cu + build-cuda-fp32: runs-on: ubuntu-latest container: diff --git a/Makefile b/Makefile index eddbca8..9ea6866 100644 --- a/Makefile +++ b/Makefile @@ -13,7 +13,7 @@ CUDA_OUTPUT_FILE = -o $@ # NVCC flags # -t=0 is short for --threads, 0 = number of CPUs on the machine -NVCC_FLAGS = -O3 -t=0 --use_fast_math +NVCC_FLAGS = -O3 -t=0 --use_fast_math -std=c++17 NVCC_LDFLAGS = -lcublas -lcublasLt NVCC_INCLUDES = NVCC_LDLIBS = diff --git a/dev/cuda/Makefile b/dev/cuda/Makefile index 14eae20..4a14ac4 100644 --- a/dev/cuda/Makefile +++ b/dev/cuda/Makefile @@ -10,7 +10,7 @@ endif # Compiler flags CFLAGS = -O3 --use_fast_math -NVCCFLAGS = -lcublas -lcublasLt +NVCCFLAGS = -lcublas -lcublasLt -std=c++17 MPI_PATHS = -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux-gnu/openmpi/lib/ # Default rule for our CUDA files diff --git a/dev/cuda/matmul_backward.cu b/dev/cuda/matmul_backward.cu index dece1f6..6f4222b 100644 --- a/dev/cuda/matmul_backward.cu +++ b/dev/cuda/matmul_backward.cu @@ -268,12 +268,14 @@ int main(int argc, char **argv) { free(dout); free(inp); free(weight); + free(ones); cudaCheck(cudaFree(d_dinp)); cudaCheck(cudaFree(d_dweight)); cudaCheck(cudaFree(d_dbias)); cudaCheck(cudaFree(d_dout)); cudaCheck(cudaFree(d_inp)); cudaCheck(cudaFree(d_weight)); + cudaCheck(cudaFree(d_ones)); cublasCheck(cublasDestroy(cublas_handle)); return 0; diff --git a/dev/cuda/matmul_backward_bias.cu b/dev/cuda/matmul_backward_bias.cu index 16172bc..ee75697 100644 --- a/dev/cuda/matmul_backward_bias.cu +++ b/dev/cuda/matmul_backward_bias.cu @@ -92,7 +92,7 @@ __global__ void matmul_backward_bias_kernel1(floatX* dbias, const floatX* dout, } // write the final result (at thread 0) to global memory if (tid == 0) { - dbias[o] = (float)dbias[o] + shared[0]; + dbias[o] = (floatX)((float)dbias[o] + shared[0]); } } @@ -116,7 +116,7 @@ __global__ void matmul_backward_bias_kernel2(floatX* dbias, const floatX* dout, sum = cg::reduce(warp, sum, cg::plus{}); // write the result to output (global memory) if(warp.thread_rank() == 0) { - dbias[idx] += sum; + dbias[idx] += (floatX)sum; } } @@ -132,12 +132,13 @@ __global__ void matmul_backward_bias_kernel3(floatX* dbias, const floatX* dout, int warp_id = threadIdx.x / 32; int lane_id = threadIdx.x % 32; int idx = blockIdx.x; // simply one block per row - // round 1: thread coarsening to reduce the problem size from B*T to 32 + // round 1: thread coarsening to reduce the problem size from B*T to block_size float thread_sum = 0.0f; for(int i = threadIdx.x; i < BT; i += blockDim.x) { thread_sum += (float)dout[i * OC + idx]; } // now do a warp-level reduce to get the sum across the 32 threads in each warp + // reduce the problem size from block_size to block_size/32 i.e. `num_warps` float warp_sum = cg::reduce(warp, thread_sum, cg::plus{}); // store the warp sum in shared memory (we could have lane_id == 0 guard but not needed) shared_sum[warp_id] = warp_sum; @@ -167,7 +168,7 @@ __global__ void matmul_backward_bias_kernel4(floatX* dbias, const floatX* dout, const int vstep = blockDim.x / warpSize; // number of warps in a block, e.g. 4 // 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 + // so e.g. 4 (`vstep`) threads (of the same lane_id) will reduce this one column const floatX* dout_col = dout + tl + lane_id; // column reductions by looping through the rows @@ -661,6 +662,7 @@ int main(int argc, char **argv) { // cleanups free(dbias); free(dout); + cudaCheck(cudaFree(dbias_buffer)); cudaCheck(cudaFree(d_dbias)); cudaCheck(cudaFree(d_dout)); diff --git a/dev/cuda/residual_forward.cu b/dev/cuda/residual_forward.cu index fd7d1fb..0f3fd7e 100644 --- a/dev/cuda/residual_forward.cu +++ b/dev/cuda/residual_forward.cu @@ -99,7 +99,7 @@ int main(int argc, char **argv) { float* out = (float*)malloc(B * T * C * sizeof(float)); float* inp1 = make_random_float(B * T * C); float* inp2 = make_random_float(B * T * C); - + // move to GPU floatX* d_out; floatX* d_inp1; diff --git a/dev/cuda/softmax_forward.cu b/dev/cuda/softmax_forward.cu index e6b4616..3fb57fb 100644 --- a/dev/cuda/softmax_forward.cu +++ b/dev/cuda/softmax_forward.cu @@ -135,7 +135,6 @@ __global__ void softmax_forward_kernel2(float* out, const float* inp, int N, int maxval = fmaxf(maxval, x[i]); } shared[tid] = maxval; - __syncthreads(); // reductions for (int stride = block_size / 2; stride >= 1; stride /= 2) { __syncthreads(); @@ -157,7 +156,6 @@ __global__ void softmax_forward_kernel2(float* out, const float* inp, int N, int sumval += x[i]; } shared[tid] = sumval; - __syncthreads(); // reductions for (int stride = block_size / 2; stride >= 1; stride /= 2) { __syncthreads(); @@ -210,14 +208,13 @@ __global__ void softmax_forward_kernel3(float* out, const float* inp, int N, int for (int i = tid; i < C; i += blockDim.x) { sumval += x[i]; } + // No need to broadcast sumval since all threads in the warp will have the same value + // (due to the fact that we're using __shfl_xor_sync) sumval = warpReduceSum(sumval); - // Broadcast sumval within the warp - float sum = __shfl_sync(0xFFFFFFFF, sumval, 0); - // Divide the input values by the sum for (int i = tid; i < C; i += blockDim.x) { - out[idx * C + i] = x[i] / sum; + out[idx * C + i] = x[i] / sumval; } } @@ -238,10 +235,9 @@ __global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int // the number of warps per block. recall that blockDim.x is block_size int warpsPerBlock = blockDim.x / 32; - // shared[] must be allocated to have 2 * warpsPerBlock elements - // first half for max values, the second half for sum values - float* maxvals = shared; - float* sumvals = &shared[warpsPerBlock]; + // shared[] must be allocated to have warpsPerBlock elements + // those will be used for max and sum values + float* max_or_sum_storage = shared; // one row of inp, i.e. inp[idx, :] of shape (C,) const float* x = inp + idx * C; @@ -255,21 +251,21 @@ __global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int maxval = warpReduceMax(maxval); // the 0th thread of each warp writes the maxval of that warp to shared memory - if (laneId == 0) maxvals[warpId] = maxval; + if (laneId == 0) max_or_sum_storage[warpId] = maxval; __syncthreads(); - // now the 0th thread reduces the maxvals in shared memory, i.e. across warps + // now the 0th thread of the block reduces the max values in shared memory, i.e. across warps if (tid == 0) { - float val = maxvals[tid]; + float val = max_or_sum_storage[tid]; for (int i = 1; i < warpsPerBlock; i++) { - val = fmaxf(val, maxvals[i]); + val = fmaxf(val, max_or_sum_storage[i]); } // store the final max in the first position - maxvals[0] = val; + max_or_sum_storage[0] = val; } __syncthreads(); // broadcast the max to all threads - float offset = maxvals[0]; + float offset = max_or_sum_storage[0]; // compute expf and write the result to global memory for (int i = tid; i < C; i += blockDim.x) { @@ -289,20 +285,20 @@ __global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int sumval = warpReduceSum(sumval); // write sumval to shared memory - if (laneId == 0) sumvals[warpId] = sumval; + if (laneId == 0) max_or_sum_storage[warpId] = sumval; __syncthreads(); // inter-thread reduction of sum if (tid == 0) { - float val = sumvals[tid]; + float val = max_or_sum_storage[tid]; for (int i = 1; i < warpsPerBlock; ++i) { - val += sumvals[i]; + val += max_or_sum_storage[i]; } - sumvals[0] = val; + max_or_sum_storage[0] = val; } __syncthreads(); // broadcast the sum to all threads - float sum = sumvals[0]; + float sum = max_or_sum_storage[0]; // divide the whole row by the sum for (int i = tid; i < C; i += blockDim.x) { @@ -322,12 +318,13 @@ __global__ void softmax_forward_online_kernel1(float* out, const float* inp, int double sum = 0.0; for (int j = 0; j < C; j++) { float maxval_prev = maxval; - if (inp_row[j] > maxval) { - maxval = inp_row[j]; - sum = sum * expf(maxval_prev - maxval) + expf(inp_row[j] - maxval); + float current_val = inp_row[j]; + if (current_val > maxval) { + maxval = current_val; + sum = sum * expf(maxval_prev - maxval) + expf(current_val - maxval); } else { - sum += expf(inp_row[j] - maxval); + sum += expf(current_val - maxval); } } @@ -590,7 +587,8 @@ void softmax_forward3(float* out, const float* inp, int N, int C, int block_size void softmax_forward4(float* out, const float* inp, int N, int C, int block_size) { int grid_size = N; - size_t shared_mem_size = 2 * block_size / 32 * sizeof(float); + // for each warp in the block we need a float that will be used for both maxval and sumval + size_t shared_mem_size = block_size / 32 * sizeof(float); softmax_forward_kernel4<<>>(out, inp, N, C); } @@ -672,11 +670,10 @@ int main(int argc, char **argv) { const int* outliers = make_random_int(B * T * 3, V); for(int k = 0; k < 3; ++k) { for(int j = 0; j < B * T; ++j) { - inp[j * V + outliers[j*3 + k]] *= 20; + inp[j * V + outliers[j*3 + k]] *= 20; } } - // move to GPU float* d_out; float* d_inp; @@ -728,6 +725,7 @@ int main(int argc, char **argv) { // free memory free(out); free(inp); + free((void*)outliers); cudaCheck(cudaFree(d_out)); cudaCheck(cudaFree(d_inp)); diff --git a/dev/cuda/trimat_forward.cu b/dev/cuda/trimat_forward.cu index 1c093e2..e05f28b 100644 --- a/dev/cuda/trimat_forward.cu +++ b/dev/cuda/trimat_forward.cu @@ -65,30 +65,38 @@ static float* d_qkvr; // scratch for the cublas kernel // taken from then attention forward pass void trimul_cpu(float* out, const float* inp, int B, int T, int C, int NH) { + // inp shape: (B, T, 3, NH, HS) + // out shape: (B, NH, T, T) int C3 = C*3; - int hs = C / NH; // head size - float scale = 1.0 / sqrtf(hs); + int HS = C / NH; // head size + float scale = 1.0 / sqrtf(HS); for (int b = 0; b < B; b++) { for (int t = 0; t < T; t++) { - for (int h = 0; h < NH; h++) { - const float* query_t = inp + b * T * C3 + t * C3 + h * hs; - float* out_bth = out + b * NH * T * T + h * T * T + t * T; + for (int nh = 0; nh < NH; nh++) { + // Q[b][nh][t][:] = inp[b][t][0][nh][:] (where : is the slice operator for hs) + const float* query_t = inp + b * T * C3 + t * C3 + nh * HS; + // out[b][nh][t][:] + float* out_bth = out + b * NH * T * T + nh * T * T + t * T; // pass 1: calculate query dot key and maxval for (int t2 = 0; t2 <= t; t2++) { - const float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key + // K[b][nh][t2][:] = inp[b][t2][1][nh][:] + const float* key_t2 = inp + b * T * C3 + t2 * C3 + nh * HS + C; // +C because it's key - // (query_t) dot (key_t2) + // Q[b][nh][t][:] dot K[b][nh][t2][:] float val = 0.0f; - for (int i = 0; i < hs; i++) { + for (int i = 0; i < HS; i++) { val += query_t[i] * key_t2[i]; } val *= scale; + // out[b][nh][t][t2] = val out_bth[t2] = val; } for(int t2 = t + 1; t2 < T; ++t2) { + // causal mask, using NAN to supress warnings -> it could be -inf + // but it doesn't matter because in validate_result we ignore infinities/NANs out_bth[t2] = NAN; } } @@ -98,31 +106,31 @@ void trimul_cpu(float* out, const float* inp, __global__ void permute_kernel(float* q, float* k, float* v, const float* inp, - int B, int N, int NH, int d) { - // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, N, d) - // but instead, we have a single tensor QKV (inp) of shape (B, N, 3, NH, d) + int B, int T, int NH, int HS) { + // okay so now, this kernel wants Q,K,V to all be of shape (B, NH, T, HS) + // but instead, we have a single tensor QKV (inp) of shape (B, T, 3, NH, HS) int idx = blockIdx.x * blockDim.x + threadIdx.x; - // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_] + // Q[b][nh][t][hs] = inp[b][t][0][nh][hs] - if (idx < B * NH * N * d) { - int b = idx / (NH * N * d); - int rest = idx % (NH * N * d); - int nh_ = rest / (N * d); - rest = rest % (N * d); - int n = rest / d; - int d_ = rest % d; + if (idx < B * NH * T * HS) { + int b = idx / (NH * T * HS); + int rest = idx % (NH * T * HS); + int nh = rest / (T * HS); + rest = rest % (T * HS); + int t = rest / HS; + int hs = rest % HS; int inp_idx = \ - (b * N * 3 * NH * d) - + (n * 3 * NH * d) - + (0 * NH * d) - + (nh_ * d) - + d_; + (b * T * 3 * NH * HS) + + (t * 3 * NH * HS) + + (0 * NH * HS) + + (nh * HS) + + hs; q[idx] = inp[inp_idx]; - k[idx] = inp[inp_idx + NH * d]; - v[idx] = inp[inp_idx + 2 * (NH * d)]; + k[idx] = inp[inp_idx + NH * HS]; + v[idx] = inp[inp_idx + 2 * (NH * HS)]; } } @@ -145,6 +153,35 @@ void trimul_cublas(float* preatt, // batched matrix multiply with cuBLAS const float alpha = 1.0f / sqrtf(HS); const float beta = 0.0f; + // This schedules in parallel B*NH matmuls of shape q@k^t = (T, HS) @ (HS, T) = (T, T). + // IMPORTANT NOTE: Cublas uses a column-major (and we use row-major in our codebase) representation, + // so this call might look confusing to you if you look at the `cublasSgemmStridedBatched` signature. + // + // In order to avoid having to do an additional transpose operation after this func call, + // we need to pass in K as the first argument and Q as the second argument, which might make you think we're computing K^T @ Q. + // That combined with the shapes we got after the permute kernel - (B, NH, T, HS) (I'll omit B, NH for brevity going forward) + // and you might think we end up with (HS, T) @ (T, HS) = (HS, HS). + // This is not the case. :) + // + // Cublas sees our row-major matrix (T, HS) as (HS, T), hence we set the lead dimensions to HS (see function signature). + // We transpose K and end up computing K^T @ Q = (T, HS) @ (HS, T) = (T, T). + // If you were to interpret the above formula K^T @ Q you might think we end up with: + // ----------------------------------- + // k1.dot(q1) k1.dot(q2) ... k1.dot(qT) + // k2.dot(q1) k2.dot(q2) ... k2.dot(qT) + // ... + // kT.dot(q1) kT.dot(q2) ... kT.dot(qT) + // ----------------------------------- + // But as I mentioned, Cublas is column-major! + // So given that the dot product is symmetric we can write k1.dot(q1) as q1.dot(k1) and transposing the above + // representation we can see what we actually end up with in the row-major format: + // ----------------------------------- + // q1.dot(k1) q1.dot(k2) ... q1.dot(kT) + // q2.dot(k1) q2.dot(k2) ... q2.dot(kT) + // ... + // qT.dot(k1) qT.dot(k2) ... qT.dot(kT) + // ----------------------------------- + // which is exactly what we wanted! :) cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, @@ -173,7 +210,7 @@ void trimul_cublas(float* preatt, */ // using creates an alias for a function pointer -using matmul_fn_ptr = void(*)(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha); +using matmul_fn_ptr = void(*)(float* p, int PS, const float* k, int KS, const float* q, int QS, int T, int HS, float alpha); template __global__ void __launch_bounds__(256, 2) trimul_global(float* out, const float* inp, int T, int C, int NH) { @@ -183,20 +220,21 @@ __global__ void __launch_bounds__(256, 2) trimul_global(float* out, const float* // set up indices int C3 = C*3; - int hs = C / NH; // head size - float scale = 1.0 / sqrtf(hs); + int HS = C / NH; // head size + float scale = 1.0 / sqrtf(HS); // we put the "batch x head" dimension into the z block index. - int h = blockIdx.z % NH; int b = blockIdx.z / NH; + int nh = blockIdx.z % NH; // Get the base address for the current batch and head - const float* q = inp + b * T * C3 + h * hs; - const float* k = inp + b * T * C3 + h * hs + C; - float* r = out + (b*NH + h)*T*T; + // shapes -> inp (B, T, 3, NH, HS), Q (B, NH, T, HS), K (B, NH, T, HS) + const float* q = inp + b * T * C3 + nh * HS; // Q[b][nh][:][:] = inp[b][:][0][nh][:] + const float* k = inp + b * T * C3 + nh * HS + C; // K[b][nh][:][:] = inp[b][:][1][nh][:] + float* r = out + (b*NH + nh)*T*T; // out[b][nh][:][:] // start the multiplication - matmul_tri(r, T, q, C3, k, C3, T, hs, scale); + matmul_tri(r, T, k, C3, q, C3, T, HS, scale); } template @@ -239,12 +277,22 @@ void trimul_launcher(float* out, const float* inp, int B, int T, int C, int NH) */ // baseline implementation: 20 ms -__device__ void matmul_tri_naive(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { - // get coordinates of our block +__device__ void matmul_tri_naive(float* p, int PS, const float* k, int KS, const float* q, int QS, int T, int HS, float alpha) { + // coordinate system: + // | - - - - - > j + // | + // | + // v + // i + // get coordinates of our block - each thread is responsible for a single 8x8 block. int i_base = 128 * blockIdx.x + 8 * threadIdx.x; int j_base = 128 * blockIdx.y + 8 * threadIdx.y; - // one more check to skip the upper diagonal in blocks that are on the diagonal. + // One more check to skip the upper diagonal in blocks that are on the diagonal. + // Note: we deliberately waste some compute on the jagged diagonal i.e. elements that belong + // to the upper triangle that should be masked out. This will be ignored due to the causal mask + // in the reference CPU implementation when used in the `validate_result` function. + // Alternatively this check should be done in the nested for loop below -> if (i > j) return. if(j_base > i_base) return; @@ -254,17 +302,17 @@ __device__ void matmul_tri_naive(float* p, int ps, const float* k, int ks, const for(int jo = 0; jo < 8; ++jo) { int j = j_base + jo; float val = 0; - for (int s = 0; s < hs; ++s) { - val += k[i * ks + s] * q[j * qs + s]; + for (int s = 0; s < HS; ++s) { + val += q[i * QS + s] * k[j * KS + s]; } - p[i * ps + j] = val * alpha; + p[i * PS + j] = val * alpha; } } } /* ** Chapter IV - ... ** * - * Each worker is producing 64 combined cookies from 8 animals and 8 landscapes. They send there runners of 64 times + * Each worker is producing 64 combined cookies from 8 animals and 8 landscapes. They send their runners 64 times * to fetch the corresponding shapes. This is terribly inefficient; The runners need a minute or so for each trip, * but making a cookie can be done in just a second. * @@ -292,7 +340,7 @@ __device__ void matmul_tri_naive(float* p, int ps, const float* k, int ks, const */ // reorganize loops to enable data reuse: 3.5 ms -__device__ void matmul_tri_registers(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { +__device__ void matmul_tri_registers(float* p, int PS, const float* k, int KS, const float* q, int QS, int T, int HS, float alpha) { int i_base = 128 * blockIdx.x + 8 * threadIdx.x; int j_base = 128 * blockIdx.y + 8 * threadIdx.y; @@ -300,17 +348,17 @@ __device__ void matmul_tri_registers(float* p, int ps, const float* k, int ks, c return; // shift our pointers to the sub-block this thread is responsible for - k += i_base * ks; - q += j_base * qs; - p += i_base * ps + j_base; + q += i_base * QS; + k += j_base * KS; + p += i_base * PS + j_base; float vals[8][8] = {}; - for (int s = 0; s < hs; ++s) { + for (int hs = 0; hs < HS; ++hs) { float lhs[8]; float rhs[8]; for (int u = 0; u < 8; ++u) { - lhs[u] = k[u * ks + s]; - rhs[u] = q[u * qs + s]; + lhs[u] = q[u * QS + hs]; + rhs[u] = k[u * KS + hs]; } for (int i = 0; i < 8; ++i) { @@ -322,7 +370,7 @@ __device__ void matmul_tri_registers(float* p, int ps, const float* k, int ks, c for (int i = 0; i < 8; ++i) { for (int j = 0; j < 8; ++j) { - p[i * ps + j] = vals[i][j] * alpha; + p[i * PS + j] = vals[i][j] * alpha; } } } @@ -334,7 +382,7 @@ __device__ void matmul_tri_registers(float* p, int ps, const float* k, int ks, c * "Of course", the runner answers, "but they've asked me for an elephant, a lion, a zebra, and a goldfish. These * are all over the place, I can't just pick them up at one spot (_strided acccess_). * "But the lion is right next to the palm tree. You could bring those two together?", you confirm. - * "Yes", he says, "if the just asked for the different categories at the same time, that would make things + * "Yes", he says, "if they just asked for the different categories at the same time, that would make things * so much easier. See, I have this bucket, I could carry lots of things in one go if I could just scoop them up * from the same place (_coalesced access_). * @@ -364,7 +412,8 @@ __device__ void st_vec(float* address, float4 val) { } // vector instructions for coalesced memory access: 1.7 ms -__device__ void matmul_tri3(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { +__device__ void matmul_tri3(float* p, int PS, const float* k, int KS, const float* q, int QS, int T, int HS, float alpha) { + // Same logic as previous kernel we just load in float4 to improve coalescing int i_base = 128 * blockIdx.x + 8 * threadIdx.x; int j_base = 128 * blockIdx.y + 8 * threadIdx.y; @@ -372,21 +421,21 @@ __device__ void matmul_tri3(float* p, int ps, const float* k, int ks, const floa return; // shift our pointers to the sub-block this thread is responsible for - k += i_base * ks; - q += j_base * qs; - p += i_base * ps + j_base; + q += i_base * QS; + k += j_base * KS; + p += i_base * PS + j_base; float vals[8][8] = {}; - for (int s = 0; s < hs; s += 4) { + for (int hs = 0; hs < HS; hs += 4) { // load in float4 to improve coalescing float4 rhs[8]; for (int u = 0; u < 8; ++u) { - rhs[u] = ld_vec(q + u * qs + s); + rhs[u] = ld_vec(k + u * KS + hs); } for (int i = 0; i < 8; ++i) { // no need to keep lhs around for the i loop, its only reused in the j loop anyway. - float4 lhs = ld_vec(k + i * ks + s); + float4 lhs = ld_vec(q + i * QS + hs); for (int j = 0; j < 8; ++j) { vals[i][j] += lhs.x * rhs[j].x; vals[i][j] += lhs.y * rhs[j].y; @@ -403,7 +452,7 @@ __device__ void matmul_tri3(float* p, int ps, const float* k, int ks, const floa result.y = vals[i][j + 1] * alpha; result.z = vals[i][j + 2] * alpha; result.w = vals[i][j + 3] * alpha; - st_vec(p + i * ps + j, result); + st_vec(p + i * PS + j, result); } } } @@ -424,7 +473,7 @@ __device__ void matmul_tri3(float* p, int ps, const float* k, int ks, const floa * details.] * */ -__device__ void matmul_tri4(float* p, int ps, const float* k, int ks, const float* q, int qs, int T, int hs, float alpha) { +__device__ void matmul_tri4(float* p, int PS, const float* k, int KS, const float* q, int QS, int T, int HS, float alpha) { int i_base = 128 * blockIdx.x + 8 * threadIdx.x; int j_base = 128 * blockIdx.y + 8 * threadIdx.y; @@ -433,14 +482,14 @@ __device__ void matmul_tri4(float* p, int ps, const float* k, int ks, const floa if (blockIdx.y > blockIdx.x) return; - k += 128 * blockIdx.x * ks; - q += 128 * blockIdx.y * qs; + q += 128 * blockIdx.x * QS; + k += 128 * blockIdx.y * KS; __shared__ float lhs_s[128][32]; __shared__ float rhs_s[128][32]; float vals[8][8] = {}; - for (int so = 0; so < hs; so += 32) { + for (int so = 0; so < HS; so += 32) { // Read a large slice of the input, worked on together by all threads. // They are organized differently for this part. We want to ensure // fully coalesced loads, so we let a single warp handle consecutive @@ -448,14 +497,23 @@ __device__ void matmul_tri4(float* p, int ps, const float* k, int ks, const floa // in one read operation. // note: threads may read data here that they don't need themselves. // this really is a block-level operation. + // note2: 16x16 threads (i.e. the block) will, through this for loop, fetch 32 dims from 128 keys and 128 queries + // i.e. from Q/K, of shape (T, HS) take q[:128, so*32:(so+1)*32] and k[:128, so*32:(so+1)*32] __syncthreads(); for(int y = threadIdx.y / 2; y < 128; y += 8) { int xo = (threadIdx.y % 2) * 16; - lhs_s[y][threadIdx.x + xo] = k[y * ks + so + threadIdx.x + xo]; - rhs_s[y][threadIdx.x + xo] = q[y * qs + so + threadIdx.x + xo]; + lhs_s[y][threadIdx.x + xo] = q[y * QS + so + threadIdx.x + xo]; + rhs_s[y][threadIdx.x + xo] = k[y * KS + so + threadIdx.x + xo]; } __syncthreads(); + // Now we compute a partial dot product (only 32 dims) for all combinations of keys and queries (128x128). + // Each thread does 8x8 of these partial dot products. + // E.g. thread (0,0) covers queries 0-7 and keys 0-7. More generally first row of threads + // (0,:) covers queries 0-7 with keys 0-127 and so on. + // In the next iterations of the outer (`so`) loop we'll be accumulating values to `vals` until we + // get the full dot product. We then later deposit it into the output matrix for all 8x8 blocks + // that are below the diagonal. for (int si = 0; si < 32; ++si) { float rhs[8]; for (int u = 0; u < 8; ++u) { @@ -484,7 +542,7 @@ __device__ void matmul_tri4(float* p, int ps, const float* k, int ks, const floa result.y = vals[ii][ji + 1] * alpha; result.z = vals[ii][ji + 2] * alpha; result.w = vals[ii][ji + 3] * alpha; - st_vec(p + i * ps + j, result); + st_vec(p + i * PS + j, result); } } } diff --git a/train_gpt2.cu b/train_gpt2.cu index f9eeac3..c28aca5 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -1063,20 +1063,22 @@ float multi_gpu_cpu_float_sum(float value) { // Averages out the loss and gradients across all GPUs. No-op when multi-GPU is disabled. // todo - this version only works if all the parameters are the same size (floatX) -void gpt2_multi_gpu_grad_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) { +void gpt2_multi_gpu_loss_and_grad_reduce(GPT2* model, MultiGpuConfig* multi_gpu_config) { #ifdef MULTI_GPU NVTX_RANGE_FN(); + // If there's only one process, there is nothing to do if (multi_gpu_config->num_processes == 1) { return; } // Average all losses. model->accumulated_mean_loss = multi_gpu_cpu_float_sum(model->mean_loss) / multi_gpu_config->num_processes; + // Now average the gradients if(multi_gpu_config->zero_stage == 0) { - // no ZERO == standard DDP: Average all gradients. + // no ZERO == standard DDP: Average all gradients. ncclCheck(ncclAllReduce(model->grads_memory, model->grads_memory, model->num_parameters, ncclFloatX, ncclAvg, multi_gpu_config->nccl_comm, 0)); } else if (multi_gpu_config->zero_stage == 1) { - // ZERO-1: Get average gradient for local shard + // ZERO-1: Get the average gradient only for local shard floatX* local_grads_memory = (floatX*) model->grads_memory + multi_gpu_config->shard_offset; ncclCheck(ncclReduceScatter(model->grads_memory, local_grads_memory, multi_gpu_config->shard_num_parameters, @@ -1120,12 +1122,12 @@ float gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, fl // repurposing this buffer (which isn't needed now) to write grad norm into it float* grad_norm_squared = (float*)model->acts.output; if (multi_gpu_config->zero_stage == 1) { - // ^1 because of the ncclReduceScatter() in gpt2_multi_gpu_grad_reduce, + // ^1 because of the ncclReduceScatter() in gpt2_multi_gpu_loss_and_grad_reduce, // grads_memory only contains the averaged gradients at the local shard // so we only calculate the grad norm at the grads_memory belonging to the local shard global_norm_squared(grad_norm_squared, grads_memory + shard_offset, shard_num_parameters); } else { - // the ncclAllReduce() in gpt2_multi_gpu_grad_reduce has averaged the gradients across all GPUs + // the ncclAllReduce() in gpt2_multi_gpu_loss_and_grad_reduce has averaged the gradients across all GPUs // so each GPU can compute the squared norm over the whole grad vector, with no added comms needed global_norm_squared(grad_norm_squared, grads_memory, model->num_parameters); } @@ -1228,15 +1230,23 @@ void gpt2_multi_gpu_param_gather(GPT2 *model, MultiGpuConfig* multi_gpu_config) } float gpt2_estimate_mfu(GPT2 *model, int num_tokens, float dt) { - // estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS - // see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311 - // TODO this calculation is only valid for an A100: generalize it? - int N = model->num_parameters; + /* + Estimate model flops utilization (MFU) + ref: Section 2.1 of https://arxiv.org/pdf/2001.08361 + Note: Ideally, the N here would be only the parameters that actually + participate in matrix multiplications. In this N, we are over-estimating by + including LayerNorm params, biases, and the position embedding weights, + but these are very small terms. Also keep in mind that we would want to exclude + the token embedding weights, but in GPT-2 these are weight shared, so they + participate in the classifier matmul, so they are correct to be included in N. + Note 2: The first term (6 * N) in flops_per_token is all weight matmuls, the + second is the attention matmul, which is also usually a small contribution. + */ + size_t N = model->num_parameters; int L = model->config.num_layers; - int H = model->config.num_heads; - int Q = model->config.channels / model->config.num_heads; + int C = model->config.channels; int T = model->seq_len; - size_t flops_per_token = (size_t)6 * N + (size_t)12 * L * H * Q * T; + size_t flops_per_token = 6 * N + (size_t)6 * L * C * T; size_t flops_per_step = flops_per_token * num_tokens; // express our flops throughput as ratio of A100 bfloat16 peak flops float flops_achieved = (float)flops_per_step * (1.0f / dt); // per second @@ -1789,8 +1799,8 @@ int main(int argc, char *argv[]) { // override the mean loss, accounting for the gradient accumulation loop // this is esp important to do here in multigpu update below, where model.mean_loss gets allreduced model.mean_loss = lossf; - // update the parameters - gpt2_multi_gpu_grad_reduce(&model, &multi_gpu_config); + // average the loss and the gradients between all processes + gpt2_multi_gpu_loss_and_grad_reduce(&model, &multi_gpu_config); // learning rate schedule: warmup linearly to max LR, then cosine decay to LR * final_learning_rate_frac float step_learning_rate = learning_rate; if (step < warmup_iterations) { diff --git a/train_gpt2.py b/train_gpt2.py index 5c0c0e5..1ad6d35 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -32,6 +32,7 @@ import torch._inductor.config as config from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed import init_process_group, destroy_process_group from torch.distributed.optim import ZeroRedundancyOptimizer +import torch.distributed as dist # ----------------------------------------------------------------------------- # PyTorch nn.Module definitions for the GPT-2 model @@ -812,7 +813,7 @@ if __name__ == "__main__": # addition of gradients corresponds to a SUM in the objective, but # instead of a SUM we want MEAN, so we scale the loss here loss = loss / grad_accum_steps - lossf += loss.item() # keep track of the mean loss + lossf += loss.detach() # keep track of the mean loss # backward pass if ddp: # we want only the last micro-step to sync grads in a DDP model @@ -821,6 +822,9 @@ if __name__ == "__main__": model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1) if not args.inference_only: loss.backward() + if ddp: + dist.all_reduce(lossf, op=dist.ReduceOp.AVG) + lossf = lossf.item() norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) # determine and set the learning rate for this iteration lr = get_lr(step)