mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-27 20:25:09 -04:00
Merge branch 'karpathy:master' into master
This commit is contained in:
Vyom Sharma
GitHub
commit
c24ef4a4b4
10 changed files with 217 additions and 115 deletions
28
.github/workflows/ci.yml
vendored
28
.github/workflows/ci.yml
vendored
|
|
@ -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:
|
||||
|
|
|
|||
2
Makefile
2
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 =
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<float>{});
|
||||
// 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<float>{});
|
||||
// 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));
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<<<grid_size, block_size, shared_mem_size>>>(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));
|
||||
|
||||
|
|
|
|||
|
|
@ -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<matmul_fn_ptr matmul_tri>
|
||||
__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<matmul_fn_ptr matmul_tri>
|
||||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue