diff --git a/TESTING_GUIDE.md b/TESTING_GUIDE.md index 6e83ad2..00fcc33 100644 --- a/TESTING_GUIDE.md +++ b/TESTING_GUIDE.md @@ -1,450 +1,64 @@ -# Complete Installation, Testing & Verification Guide -## Prerequisites Check +# Quick Testing Guide (concise) -### Hardware Requirements +This file contains the minimal steps to get a 2-node NVSHMEM/MPI training run working for the `pipeline_llm.c` project. + +**Requirements:** CUDA-enabled Linux machines with MPI and NVSHMEM installed; 2 GPUs for the basic pipeline test. + +**1) Minimal environment** + +- **Set NVSHMEM path and runtime vars:** ```bash -# Check GPU count (need exactly 2) -nvidia-smi --query-gpu=count --format=csv,noheader - -# Check GPU models (A100, V100, or compute capability 7.0+) -nvidia-smi --query-gpu=name,compute_cap --format=csv - -# Expected output: -# NVIDIA A100-SXM4-40GB, 8.0 -# NVIDIA A100-SXM4-40GB, 8.0 -``` - -### Software Requirements - -```bash -# Check CUDA version (need 11.0+) -nvcc --version - -# Check MPI -mpirun --version - -# Check if on Linux -uname -a # Should show Linux, not Darwin (Mac) or Windows -``` - ---- - -## Step 1: Install NVSHMEM - -### Download and Install - -```bash -# Go to scratch directory -cd /scratch/$USER - -# Download NVSHMEM (adjust version for your CUDA) -wget https://developer.download.nvidia.com/compute/nvshmem/redist/libnvshmem/linux-x86_64/libnvshmem-linux-x86_64-3.4.5_cuda12-archive.tar.xz - -# Extract -tar -xf libnvshmem-linux-x86_64-3.4.5_cuda12-archive.tar.xz - -# Set environment variables (add to ~/.bashrc for persistence) -export NVSHMEM_HOME=/scratch/$USER/libnvshmem-linux-x86_64-3.4.5_cuda12-archive +export NVSHMEM_HOME=/path/to/libnvshmem export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH +export NVSHMEM_BOOTSTRAP=MPI ``` -### Verify Installation +**2) Build** ```bash -ls $NVSHMEM_HOME -# Should show: bin/ include/ lib/ - -ls $NVSHMEM_HOME/lib/libnvshmem* -# Should show several .so files -``` - ---- - -## Step 2: Get Data and Model - -### Download GPT-2 Model Checkpoint - -```bash -cd /scratch/$USER/pipeline_llm.c - -# Download GPT-2 124M checkpoint (~500 MB) -wget https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/gpt2_124M.bin - -# Download tokenizer -wget https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/gpt2_tokenizer.bin - -# Verify downloads -ls -lh gpt2_124M.bin gpt2_tokenizer.bin -# gpt2_124M.bin should be ~475 MB -# gpt2_tokenizer.bin should be ~0.5 MB -``` - -### Download Training Data - -```bash -# Create data directory -mkdir -p dev/data/tinyshakespeare - -# Download pre-tokenized tiny shakespeare -cd dev/data/tinyshakespeare -wget https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/tinyshakespeare_train.bin -O tiny_shakespeare_train.bin -wget https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/tinyshakespeare_val.bin -O tiny_shakespeare_val.bin - -# Go back to project root -cd ../../../ - -# Verify data files -ls -lh dev/data/tinyshakespeare/ -# Should show tiny_shakespeare_train.bin (~1 MB) and tiny_shakespeare_val.bin -``` - ---- - -## Step 3: Compile the Code - -### Set Environment Variables - -```bash -# Required for compilation -export NVSHMEM_HOME=/scratch/$USER/libnvshmem-linux-x86_64-3.4.5_cuda12-archive -export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH -``` - -### Compile - -```bash -cd /scratch/$USER/pipeline_llm.c - -# Build with Makefile +cd /path/to/pipeline_llm.c make nvshmem_train_gpt2 - -# Expected output: -# nvcc --threads=0 -t=0 --use_fast_math -std=c++17 -O3 \ -# -arch=sm_80 \ -# -I/scratch/.../nvshmem/include -L/scratch/.../nvshmem/lib \ -# -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux-gnu/openmpi/lib \ -# nvshmem_train_gpt2.cu -lnvshmem -lcublas -lcublasLt -lcudart -lnvidia-ml -lmpi -o nvshmem_train_gpt2 ``` -### Troubleshooting Compilation +If your GPU compute capability differs, adjust `-arch=sm_XX` in the Makefile. -**Error: `NVSHMEM_HOME not found`** +**3) Prepare model & data (examples)** + +Place model and tokenizers under the repo or a data path you use. Example filenames used by scripts: + +- `gpt2_124M.bin` +- tokenizer file (project uses `llmc` tokenizer) + +Small datasets (e.g. tiny shakespeare) can be put under `dev/data/` for quick tests. + +**4) Run a basic 2-node test** ```bash -# Make sure NVSHMEM_HOME is set correctly -echo $NVSHMEM_HOME -# Should print: /scratch/$USER/libnvshmem-linux-x86_64-3.4.5_cuda12-archive - -# If not, export it again -export NVSHMEM_HOME=/scratch/$USER/libnvshmem-linux-x86_64-3.4.5_cuda12-archive -``` - -**Error: `cannot find -lnvshmem`** - -```bash -# Check if libraries exist -ls $NVSHMEM_HOME/lib/libnvshmem.so - -# Add to library path -export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH -``` - -**Error: `arch=sm_80 not supported`** - -```bash -# Your GPU might have different compute capability -# Check your GPU's compute capability -nvidia-smi --query-gpu=compute_cap --format=csv,noheader - -# If it shows 7.0 or 7.5 (V100), edit Makefile line 307: -# Change: -arch=sm_80 -# To: -arch=sm_70 (for V100) -``` - ---- - -## Step 4: Set Runtime Environment - -```bash -# Runtime library path -export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH - -# MPI bootstrap for NVSHMEM -export NVSHMEM_BOOTSTRAP=MPI - -# Disable InfiniBand for single-node testing -export NVSHMEM_DISABLE_IBRC=1 -export NVSHMEM_DISABLE_IBGDA=1 -export NVSHMEM_DISABLE_IBDEVX=1 -export NVSHMEM_REMOTE_TRANSPORT=none - -# Make these permanent by adding to ~/.bashrc: -cat >> ~/.bashrc << 'EOF' -# NVSHMEM Environment -export NVSHMEM_HOME=/scratch/$USER/libnvshmem-linux-x86_64-3.4.5_cuda12-archive -export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH -export NVSHMEM_BOOTSTRAP=MPI -export NVSHMEM_DISABLE_IBRC=1 -export NVSHMEM_DISABLE_IBGDA=1 -export NVSHMEM_DISABLE_IBDEVX=1 -export NVSHMEM_REMOTE_TRANSPORT=none -EOF - -source ~/.bashrc -``` - ---- - -## Step 5: Run Tests - -### Test 1: Basic Initialization (2-3 seconds) - -```bash -cd /scratch/$USER/pipeline_llm.c - -# Run with minimal settings +# Example: 2 processes, small batch and seq len for quick validation mpirun -np 2 ./nvshmem_train_gpt2 -b 1 -t 64 -v 1 -m 1 - -# Expected output: -+------------------------+----------------------------------------------------+ -| Parameter | Value | -+------------------------+----------------------------------------------------+ -| NVSHMEM PEs (GPUs) | 2 | -| This PE | 0 | -| batch size B | 1 | -| sequence length T | 64 | -+------------------------+----------------------------------------------------+ -| device | NVIDIA A100-SXM4-40GB | -+------------------------+----------------------------------------------------+ -| num_layers L | 12 | -| num_parameters | 124439808 | -+------------------------+----------------------------------------------------+ -allocated 124 MiB for model parameters -allocated 47 MiB for activations -allocated 0 MiB for NVSHMEM activation buffer -allocated 0 MiB for NVSHMEM gradient buffer ``` -**✅ Pass Criteria:** +Expected: both ranks initialize and the program runs without hangs. If it stalls, check MPI and NVSHMEM environment variables and `nvidia-smi`. -- No error messages -- Both PEs (0 and 1) initialize -- NVSHMEM buffers allocated -- Program completes without hanging - -### Test 2: Small Training Run (30 seconds - 1 minute) +**5) Quick training sanity check** ```bash -# Run 10 training steps with small batch mpirun -np 2 ./nvshmem_train_gpt2 -b 2 -t 128 -v 5 -m 5 - -# Monitor in another terminal: -watch -n 0.5 nvidia-smi - -# Expected output: -step 1/N: train loss 5.2xxx (xxx ms, xxx tok/s) -step 2/N: train loss 5.1xxx (xxx ms, xxx tok/s) -... -step 5/N: train loss 5.0xxx (xxx ms, xxx tok/s) -val loss 4.9xxx ``` -**✅ Pass Criteria:** +Look for reasonable loss values (not NaN) and decreasing trend. -- Loss values are reasonable (3.0 - 6.0 range) -- Loss decreases over steps -- Both GPUs show memory usage in `nvidia-smi` -- No deadlocks or hangs +**6) Common quick fixes** -### Test 3: Full Training Run (5-10 minutes) +- `libnvshmem.so` load error: ensure `LD_LIBRARY_PATH` contains `$NVSHMEM_HOME/lib`. +- MPI hangs: test with `mpirun -np 2 echo hello` and confirm environment variables are visible. +- CUDA OOM: reduce `-b` or `-t`. -```bash -# Run with standard settings -mpirun -np 2 ./nvshmem_train_gpt2 -b 4 -t 128 -v 10 -m 10 -s 20 +**7) Notes** -# Expected behavior: -# - Trains for full epoch on tiny shakespeare -# - Validates every 10 steps -# - Generates text samples every 20 steps -``` +- This guide is intentionally minimal — it assumes familiarity with MPI and NVSHMEM. Use the original, longer `TESTING_GUIDE.md` for full troubleshooting and advanced steps. +- For pipeline debugging start with small `B` and `T` and confirm per-rank logs. -**✅ Pass Criteria:** - -- Training completes without crashes -- Loss converges (decreases to ~3.5-4.5 range) -- Text generation improves over time -- GPU utilization ~70-90% on both GPUs - ---- - -## Step 6: Verify Correctness - -### Numerical Accuracy Test - -Compare pipeline results with single-GPU baseline: - -```bash -# Run single-GPU version (if available) -./train_gpt2_fp32cu -b 4 -t 128 | tee single_gpu.log - -# Run pipeline version -mpirun -np 2 ./nvshmem_train_gpt2 -b 4 -t 128 | tee pipeline.log - -# Compare first 10 step losses -# They should match within ~0.05 tolerance -``` - -### Memory Usage Verification - -```bash -# While training is running, check memory in another terminal: -watch -n 0.5 'nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv' - -# Expected for B=4, T=128: -# GPU 0: ~2-3 GB memory used, 70-90% utilization -# GPU 1: ~2-3 GB memory used, 70-90% utilization -``` - ---- - -## Common Issues & Solutions - -### Issue 1: `error while loading shared libraries: libnvshmem.so` - -**Solution:** - -```bash -export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH -# Make sure to do this BEFORE running mpirun -``` - -### Issue 2: `Error: This pipeline implementation requires exactly 2 GPUs, got X` - -**Solution:** - -```bash -# Check GPU visibility -nvidia-smi -CUDA_VISIBLE_DEVICES=0,1 mpirun -np 2 ./nvshmem_train_gpt2 ... -``` - -### Issue 3: Program hangs at initialization - -**Solution:** - -```bash -# Check if MPI is working -mpirun -np 2 echo "Hello from rank $OMPI_COMM_WORLD_RANK" - -# Check NVSHMEM environment -env | grep NVSHMEM - -# Try with MPI debug output -mpirun -np 2 --mca btl_base_verbose 10 ./nvshmem_train_gpt2 ... -``` - -### Issue 4: CUDA out of memory - -**Solution:** - -```bash -# Reduce batch size -mpirun -np 2 ./nvshmem_train_gpt2 -b 2 -t 64 - -# Or reduce sequence length -mpirun -np 2 ./nvshmem_train_gpt2 -b 4 -t 64 -``` - -### Issue 5: Loss is NaN or explodes - -**Solution:** - -```bash -# Check model file integrity -md5sum gpt2_124M.bin - -# Reduce learning rate -mpirun -np 2 ./nvshmem_train_gpt2 -b 4 -t 128 -l 0.0001 - -# Check for NVSHMEM communication errors in output -``` - ---- - -## Performance Expectations - -### Phase 1 (Current Implementation) - -| Metric | Expected Value | -| ---------------------- | ------------------------- | -| Pipeline Efficiency | ~50% | -| Speedup vs Single GPU | ~1.0x (no speedup) | -| GPU 0 Utilization | 70-90% during layers 0-5 | -| GPU 1 Utilization | 70-90% during layers 6-11 | -| Communication Overhead | ~5-10% of total time | - -**Why no speedup?** -Phase 1 is a **proof of concept** without micro-batching. GPUs run sequentially: - -- While GPU 0 processes layers 0-5, GPU 1 is idle -- While GPU 1 processes layers 6-11, GPU 0 is idle - -### Phase 2 (With Micro-batching - Future) - -| Metric | Target Value | -| --------------------- | ------------ | -| Pipeline Efficiency | ~85% | -| Speedup vs Single GPU | ~1.7x | -| Both GPU Utilization | ~85-95% | - ---- - -## Quick Start Commands - -### One-Liner Setup and Run - -```bash -# Set environment and run -export NVSHMEM_HOME=/scratch/$USER/libnvshmem-linux-x86_64-3.4.5_cuda12-archive && \ -export LD_LIBRARY_PATH=$NVSHMEM_HOME/lib:$LD_LIBRARY_PATH && \ -export NVSHMEM_BOOTSTRAP=MPI && \ -cd /scratch/$USER/pipeline_llm.c && \ -mpirun -np 2 ./nvshmem_train_gpt2 -b 4 -t 128 -``` - -### Monitor GPUs in Real-Time - -```bash -# Open second terminal -watch -n 0.5 nvidia-smi -``` - -### Save Training Log - -```bash -mpirun -np 2 ./nvshmem_train_gpt2 -b 4 -t 128 2>&1 | tee training_$(date +%Y%m%d_%H%M%S).log -``` - ---- - -## Next Steps After Successful Testing - -1. **Verify numerical correctness** - Compare with single-GPU baseline -2. **Measure performance metrics** - Document actual throughput and efficiency -3. **Test different batch sizes** - Find optimal B and T -4. **Profile with nsys/nvprof** - Identify bottlenecks -5. **Prepare for Phase 2** - Implement micro-batching for 1.7x speedup - ---- - -## Getting Help - -If you encounter issues not covered here: - -1. Check NVSHMEM logs: Look for NVSHMEM error messages in stderr -2. Enable MPI debugging: `mpirun --mca btl_base_verbose 10 ...` -3. Check GPU status: `nvidia-smi -l 1` (updates every second) -4. Review code: See [nvshmem_train_gpt2.cu](file:///Users/kvlnraju/Desktop/courses/semester_3/bdml/proj/pipeline_llm.c/nvshmem_train_gpt2.cu) -5. Consult docs: See [HANDOFF.md](file:///Users/kvlnraju/Desktop/courses/semester_3/bdml/proj/pipeline_llm.c/HANDOFF.md) for architecture details +**File edited:** `TESTING_GUIDE.md` — concise quick-start and troubleshooting. diff --git a/pipeline_gpt2.cu b/pipeline_gpt2.cu deleted file mode 100644 index df412ea..0000000 --- a/pipeline_gpt2.cu +++ /dev/null @@ -1,1755 +0,0 @@ -/* -GPT-2 Transformer Neural Net trained in raw CUDA -Non-trivial notes to be aware of: - -We are being clever in the backward pass to conserve memory. -In particular, all parameters use a += in the backward pass, so we -can later do gradient accumulation. But all activations have = instead of += -because these are faster (just read, no write). This is okay for all activations -except for those in the residual stream, where the gradients have to add. We make -sure that those parts work out ok and that we do a += as necessary. E.g., -the layernorms are connected to the residuals so we += in layernorm backward. -*/ - -#include -#include -#include -#include -#include -#include -#include -#include - -// GPU / CUDA related -#include -#include -#include -#include -// our own utilities -// defines: fopenCheck, freadCheck, fcloseCheck, fseekCheck, mallocCheck -#include "llmc/utils.h" -// defines: tokenizer_init, tokenizer_decode, tokenizer_free -#include "llmc/tokenizer.h" -// defines: dataloader_init, dataloader_reset, dataloader_next_batch, dataloader_free -#include "llmc/dataloader.h" - -// ---------------------------------------------------------------------------- -// CUDA utils - -// convenience macro for calculating grid/block dimensions for kernels -#define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) - -// CUDA error checking -void cudaCheck(cudaError_t error, const char *file, int line) { - if (error != cudaSuccess) { - printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, - cudaGetErrorString(error)); - exit(EXIT_FAILURE); - } -}; -#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__)) - -// cuBLAS error checking -void cublasCheck(cublasStatus_t status, const char *file, int line) -{ - if (status != CUBLAS_STATUS_SUCCESS) { - printf("[cuBLAS ERROR]: %d %s %d\n", status, file, line); - exit(EXIT_FAILURE); - } -} -#define cublasCheck(status) { cublasCheck((status), __FILE__, __LINE__); } - -static cublasComputeType_t cublas_compute_type; -cublasHandle_t cublas_handle; - -namespace cg = cooperative_groups; - -// ---------------------------------------------------------------------------- -// all the kernels - -__device__ inline float4 add_float4(const float4& a, const float4& b) { - return make_float4(a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w); -} - -// use of float4 leads to using 128-bit LDG / STG instructions in SASS, -// very helpful in memory-bound kernels like encoder_forward -__global__ void encoder_forward_kernel3(float4* out, - const int* inp, const float4* wte, const float4* wpe, - int B, int T, int C) { - int C4 = C / 4; - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int N = B * T * C4; - if (idx < N) { - int bt = idx / C4; - int b = bt / T; - int t = bt % T; - int c4 = idx % C4; - int ix = inp[b * T + t]; - out[b * T * C4 + t * C4 + c4] = add_float4(wte[ix * C4 + c4], wpe[t * C4 + c4]); - } -} - -// really bad naive kernel with atomicAdd -__global__ void encoder_backward_kernel(float* dwte, float* dwpe, - const float* dout, const int* inp, - int B, int T, int C) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int N = B * T * C; - - if (idx < N) { - int bt = idx / C; - int b = bt / T; - int t = bt % T; - int c = idx % C; - - int ix = inp[b * T + t]; - - const float* dout_btc = dout + b * T * C + t * C + c; - float* dwte_ix = dwte + ix * C + c; - float* dwpe_tc = dwpe + t * C + c; - - atomicAdd(dwte_ix, *dout_btc); - atomicAdd(dwpe_tc, *dout_btc); - } -} - -__global__ void layernorm_forward_kernel3(float* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd, - const float* __restrict__ inp, const float* __restrict__ weight, - const float* __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(); - if(idx >= N) { - return; - } - - // the row of input that this group of threads is responsible for - const float* x = inp + idx * C; - - // mean - float sum = 0.0f; - for (int i = warp.thread_rank(); i < C; i += warp.size()) { - sum += x[i]; - } - sum = cg::reduce(warp, sum, cg::plus{}); - float m = sum / C; - if(warp.thread_rank() == 0 && mean != nullptr) { - __stcs(mean + idx, m); - } - - // rstd - sum = 0.0f; - for (int i = warp.thread_rank(); i < C; i += warp.size()) { - float diff = x[i] - m; - sum += diff * diff; - } - sum = cg::reduce(warp, sum, cg::plus{}); - float s = rsqrtf(sum / C + 1e-5f); - if(warp.thread_rank() == 0 && rstd != nullptr) { - __stcs(rstd + idx, s); - } - - // final normalization and scaling by weight/bias - float* o = out + idx * C; - for (int c = warp.thread_rank(); 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 - float n = s * (__ldcs(x+c) - m); - __stcs(o+c, n * weight[c] + bias[c]); - } -} - -__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 idx = blockIdx.x * blockDim.x + threadIdx.x; - // Q[b][nh_][n][d_] = inp[b][n][0][nh_][d_] - 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; - int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; - q[idx] = __ldcs(&inp[inp_idx]); - k[idx] = __ldcs(&inp[inp_idx + NH * d]); - v[idx] = __ldcs(&inp[inp_idx + 2 * (NH * d)]); - } -} - -__global__ void permute_kernel_backward(float* dinp, - const float* dq, const float* dk, const float* dv, - int B, int N, int NH, int d) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - 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; - - int inp_idx = (b * N * 3 * NH * d) + (n * 3 * NH * d) + (0 * NH * d) + (nh_ * d) + d_; - dinp[inp_idx] = dq[idx]; - dinp[inp_idx + NH * d] = dk[idx]; - dinp[inp_idx + 2 * (NH * d)] = dv[idx]; - } -} - -__global__ void unpermute_kernel(float* inp, float *out, int B, int N, int NH, int d) { - // out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d) - int idx = blockIdx.x * blockDim.x + threadIdx.x; - // out[b][n][nh_][d_] <- inp[b][nh_][n][d_] - 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; - int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; - out[other_idx] = __ldcs(&inp[idx]); - } -} - -__global__ void unpermute_kernel_backward(float* dinp, const float *dout, int B, int N, int NH, int d) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - 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; - int other_idx = (b * NH * N * d) + (n * NH * d) + (nh_ * d) + d_; - dinp[idx] = dout[other_idx]; - } -} - -__device__ float& vec_at(float4& vec, int index) { - return reinterpret_cast(&vec)[index]; -} - -__device__ float vec_at(const float4& vec, int index) { - return reinterpret_cast(&vec)[index]; -} - -__global__ void softmax_forward_kernel5(float* out, float inv_temperature, const float* inp, int N, int T) { - // inp, out shape: (N, T, T), where N = B * NH - // fuses the multiplication by scale inside attention - // 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); - // 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 - if(idx >= N * T) { - return; - } - int own_pos = idx % T; - int pos_by_4 = own_pos / 4; - - // one row of inp, i.e. inp[idx, :] of shape (T,) - const float* x = inp + idx * T; - - // not INF, so we don't get NaNs accidentally when subtracting two values. - float maxval = -FLT_MAX; - float sumval = 0.0f; - - const float4* x_vec = reinterpret_cast(x); - for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) { - float4 v = x_vec[i]; - float old_maxval = maxval; - for(int k = 0; k < 4; ++k) { - maxval = fmaxf(maxval, vec_at(v, k)); - } - sumval *= expf(inv_temperature * (old_maxval - maxval)); - for(int k = 0; k < 4; ++k) { - sumval += expf(inv_temperature * (vec_at(v, k) - maxval)); - } - } - - if(4*pos_by_4 + warp.thread_rank() <= own_pos) { - float old_maxval = maxval; - maxval = fmaxf(maxval, x[4*pos_by_4 + warp.thread_rank()]); - sumval *= expf(inv_temperature * (old_maxval - maxval)); - sumval += expf(inv_temperature * (x[4*pos_by_4 + warp.thread_rank()] - maxval)); - } - - float global_maxval = cg::reduce(warp, maxval, cg::greater{}); - sumval *= expf(inv_temperature * (maxval - global_maxval)); - - float sum = cg::reduce(warp, sumval, cg::plus{}); - float norm = 1.f / sum; - - // divide the whole row by the sum - for (int i = warp.thread_rank(); i <= own_pos; i += warp.size()) { - // recalculation is faster than doing the round-trip through memory. - float ev = expf(inv_temperature * (__ldcs(x + i) - global_maxval)); - __stcs(out + idx * T + i, ev * norm); - } -} - -__global__ void residual_forward_kernel(float* out, float* inp1, float* inp2, int N) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < N) { - out[idx] = __ldcs(&inp1[idx]) + __ldcs(&inp2[idx]); - } -} - -#define GELU_SCALING_FACTOR sqrtf(2.0f / M_PI) -__global__ void gelu_forward_kernel(float* out, const float* inp, int N) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < N) { - float xi = inp[i]; - float cube = 0.044715f * xi * xi * xi; - out[i] = 0.5f * xi * (1.0f + tanhf(GELU_SCALING_FACTOR * (xi + cube))); - } -} - -__global__ void gelu_backward_kernel(float* dinp, const float* inp, const float* dout, const int N) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < N) { - float x = inp[i]; - float cube = 0.044715f * x * x * x; - float tanh_arg = GELU_SCALING_FACTOR * (x + cube); - float tanh_out = tanhf(tanh_arg); - float coshf_out = coshf(tanh_arg); - float sech_out = 1.0f / (coshf_out * coshf_out); - float local_grad = 0.5f * (1.0f + tanh_out) + x * 0.5f * sech_out * GELU_SCALING_FACTOR * (1.0f + 3.0f * 0.044715f * x * x); - dinp[i] = local_grad * dout[i]; - } -} - -// this kernel performs a column-wise reduction over dout, in PyTorch equivalent to: -// dbias = dout.sum((0,1)) -// the idea is to employ one block to reduce along several columns, -// where each block has a width of 32 columns to ensure coalesced access. -// at the end we accumulate the reductions performed by the warps in each block via shared memory -__global__ void matmul_backward_bias_kernel4(float* dbias, const float* dout, int B, int T, int OC) { - // this kernel is launched with 1D grid_dim of OC/32 - // for example let's say block_size is 128 - extern __shared__ float smem[]; // of size block_size (128) - const int warp_id = threadIdx.x / warpSize; // warp index in the block, 0,1,2,3 - const int lane_id = threadIdx.x % warpSize; // thread index in the warp, 0,1,2,...,31 - const int tl = blockIdx.x * warpSize; // pointer to the start column for this block - 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 - const float* dout_col = dout + tl + lane_id; - - // column reductions by looping through the rows - // each of the 4 threads offsets by its warp_id and then skips by vstep - // together these 4 threads cover all B*T rows of this (lane_id) column - // importantly, consecutive threads (in threadId) are processing adjacent columns, - // leading to a coalesced memory access pattern - float dout_sum = 0.0f; - for (int row = warp_id; row < B * T; row += vstep) { - dout_sum += dout_col[row * OC]; - } - smem[lane_id + warp_id * warpSize] = dout_sum; - __syncthreads(); - - // warp_id 0 reduces the shared memory column-wise, linearly - dout_sum = 0.0f; - if (warp_id == 0) { - for (int j = 0; j < vstep; j++) { - dout_sum += smem[lane_id + j * warpSize]; - } - dbias[tl + lane_id] += dout_sum; - } -} - -// uses shared memory instead for the reduces -__global__ void layernorm_backward_kernel2(float* dinp, float* dweight, float* dbias, - const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd, - int B, int T, int C) { - extern __shared__ float shared[]; // size = 2 * C - - namespace cg = cooperative_groups; - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); - int idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank(); - int N = B * T; - if(idx >= N) { return; } // thread guards - - int b = idx / T; - int t = idx % T; - - const float* dout_bt = dout + b * T * C + t * C; - const float* inp_bt = inp + b * T * C + t * C; - float* dinp_bt = dinp + b * T * C + t * C; - const float mean_bt = mean[b * T + t]; - const float rstd_bt = rstd[b * T + t]; - - // the first half of shared memory is bias, second is weight - float* dbias_shared = shared; - float* dweight_shared = shared + C; - - // init shared memory to zero - #pragma unroll - for(int i = threadIdx.x; i < C; i+= blockDim.x){ - dbias_shared[i] = 0.0f; - dweight_shared[i] = 0.0f; - } - __syncthreads(); - - // first: two reduce operations - float dnorm_mean = 0.0f; - float dnorm_norm_mean = 0.0f; - for (int i = warp.thread_rank(); i < C; i += warp.size()) { - float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; - float dnorm_i = weight[i] * dout_bt[i]; - dnorm_mean += dnorm_i; - dnorm_norm_mean += dnorm_i * norm_bti; - } - dnorm_mean = cg::reduce(warp, dnorm_mean, cg::plus{}); - dnorm_norm_mean = cg::reduce(warp, dnorm_norm_mean, cg::plus{}); - dnorm_mean = dnorm_mean / C; - dnorm_norm_mean = dnorm_norm_mean / C; - - // now iterate again and accumulate all the gradients - for (int i = warp.thread_rank(); i < C; i += warp.size()) { - float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt; - float dnorm_i = weight[i] * dout_bt[i]; - // gradient contribution to bias - atomicAdd(&dbias_shared[i], dout_bt[i]); - // gradient contribution to weight - atomicAdd(&dweight_shared[i], norm_bti * dout_bt[i]); - // gradient contribution to input - float dval = 0.0f; - dval += dnorm_i; // term 1 - dval -= dnorm_mean; // term 2 - dval -= norm_bti * dnorm_norm_mean; // term 3 - dval *= rstd_bt; // final scale - dinp_bt[i] += dval; - } - __syncthreads(); - - // write to global memory - for(int i = threadIdx.x; i < C; i+= blockDim.x){ - atomicAdd(&dbias[i], dbias_shared[i]); - atomicAdd(&dweight[i], dweight_shared[i]); - } -} - -__global__ void softmax_autoregressive_backward_kernel(float* dpreatt, const float* datt, const float* att, - 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 - int t0 = T - 1 - T_per_block*blockIdx.x; - - att += idx * T * T; - 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; - const float* att_bth = att + t * T; - const float* datt_bth = datt + t * T; - float* dpreatt_bth = dpreatt + t * T; - - float local_sum = 0; - for (int t2 = block.thread_rank(); t2 <= t; t2 += BlockSize) { - local_sum += att_bth[t2] * datt_bth[t2]; - } - - block_acc[warp.meta_group_rank()] = cg::reduce(warp, local_sum, cg::plus{}); - block.sync(); - local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus{}); - - for (int t3 = block.thread_rank(); 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 = __ldcs(att_bth + t3) * (__ldcs(datt_bth + t3) - local_sum); - __stcs(dpreatt_bth + t3, scale * acc); - } - } -} - -// Implements linear interpolation using only two floating-point operations (as opposed to three in a naive implementation). -// Reference: https://developer.nvidia.com/blog/lerp-faster-cuda -__device__ inline float lerp(float start, float end, float weight) { - return fma(weight, end, fma(-weight, start, start)); -} - -__global__ void adamw_kernel2(float* params_memory, float* grads_memory, float* m_memory, float* v_memory, long num_parameters, - float learning_rate, float beta1, float beta2, float beta1_correction, float beta2_correction, float eps, float weight_decay) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= num_parameters) return; // guard - float grad = grads_memory[i]; - float m = m_memory[i]; - float v = v_memory[i]; - // update the first moment (momentum) - m = lerp(grad, m, beta1); - m_memory[i] = m; - // update the second moment (RMSprop) - v = lerp(grad * grad, v, beta2); - v_memory[i] = v; - m /= beta1_correction; // m_hat - v /= beta2_correction; // v_hat - params_memory[i] -= learning_rate * (m / (sqrtf(v) + eps) + weight_decay * params_memory[i]); -} - -struct SoftmaxParams { - float Scale; - float Offset; -}; - - -__device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_tile<32>& warp, - int idx, const float* inp, int V, int P) { - // same but not float4 - // one row of inp, i.e. inp[idx, :] of shape (V,) - - const float* x = inp + idx * P; - float thread_maxval = -INFINITY; - float thread_sumval = 0.0f; - // do the loop in reverse to maximise probability of L2 cache hits - // so even small L2s get some hits on the 2nd read of the same thread - for (int i = V + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) { - float v = x[i]; - float old_maxval = thread_maxval; - thread_maxval = fmaxf(thread_maxval, v); - thread_sumval *= expf((old_maxval - thread_maxval)); - 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; - - // reduce maxval within each warp - float warp_maxval = cg::reduce(warp, thread_maxval, cg::greater{}); - // thread 0 in each warp writes to shared memory - if (lane_id == 0) { shared_maxval[warp_id] = warp_maxval; } - __syncthreads(); - // each thread now loads the maxval across previous warps - // if the thread is "out of range" of data, use -FLT_MAX as the maxval - warp_maxval = (lane_id < num_warps) ? shared_maxval[lane_id] : -FLT_MAX; - // now reduce the maxval among the warp threads - float block_maxval = cg::reduce(warp, warp_maxval, cg::greater{}); - // each thread uses maxval to scale sumval to avoid numerical instability / overflow - thread_sumval *= expf(thread_maxval - block_maxval); - // (warp-level) reduce sumval, thread 0 in each warp saves result in shared memory - float warp_sumval = cg::reduce(warp, thread_sumval, cg::plus{}); - if (lane_id == 0) { shared_sumval[warp_id] = warp_sumval; } - __syncthreads(); - // same strategy, now reduce sumval across warps - warp_sumval = (lane_id < num_warps) ? shared_sumval[lane_id] : 0.0f; - float block_sumval = cg::reduce(warp, warp_sumval, cg::plus{}); - // return the softmax parameters - return SoftmaxParams{1.f / block_sumval, block_maxval}; -} - -// same as 2 but not using float4 (see dev/cuda/classifier_fused.cu) -// will _update_ logits to logit gradients -__global__ void fused_classifier_kernel3(float* logits, float* losses, float* probs, - const float* dlosses, const int* targets, - int B, int T, int V, int P) { - namespace cg = cooperative_groups; - cg::thread_block block = cg::this_thread_block(); - cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block); - int idx = blockIdx.x; - int ix = targets[idx]; - - // softmax (reading B * T * V, same logits read again below, hopefully still in cache) - SoftmaxParams sp = prepare_softmax_blockwide_nofloat4(warp, idx, logits, V, P); - - // calculate the probability needed for the loss and update (single-threaded) - if(threadIdx.x == 0) { - float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale; - losses[idx] = -logf(prob); - } - - // very sensible default for dlosses is 1/(B*T), which is the uniform loss - float dloss = dlosses != NULL ? dlosses[idx] : 1.0f / (B*T); - // calculate the gradients directly, saves bandwidth from probs during training - // but also supports writing probs for inference-only and debugging - const float* logits_vec = logits + idx * P; - for (int i = threadIdx.x; i < V; i += blockDim.x) { - // this is the 2nd read of logits after the one in prepare_softmax2 - // this data will never be needed again, so we reduce cache persistence - float v = __ldcs(&logits_vec[i]); - float prob = expf(v - sp.Offset) * sp.Scale; - if (probs != NULL) { - probs[idx * P + i] = prob; - } - float indicator = (i == ix) ? 1.0f : 0.0f; - logits[idx * P + i] = (prob - indicator) * dloss; - } -} - -__device__ float4 ld_vec(const float* address) { - return *reinterpret_cast(address); -} - -__device__ void st_vec(float* address, float4 val) { - *reinterpret_cast(address) = val; -} - -__global__ void __launch_bounds__(16*16, 2) matmul_forward_kernel4(float* out, - const float* inp, const float* weight, const float* bias, - int C, int OC) { - // out is (B,T,OC). OC is short for "output channels", e.g. OC = 4 * C - // inp is (B,T,C), weight is (OC, C), bias is (OC) - // each thread handles 8x8 elements; each block 128 by 128 elements. - int oc = 8*(blockIdx.y * blockDim.y + threadIdx.y); - - // buffers to cache chunks of the input matrices - __shared__ float lhs_s[128][32]; - __shared__ float rhs_s[128][32]; - - // adjust our pointers for the current block - inp += 128 * blockIdx.x * C; - weight += 128 * blockIdx.y * C; - out += 128 * blockIdx.x * OC + 128 * blockIdx.y; - - float vals[8][8] = {}; - if(bias != NULL) { - for (int i = 0; i < 8; i++) { - for (int j = 0; j < 8; j += 4) { - float4 b = ld_vec(bias + oc + j); - vals[i][j+0] = b.x; - vals[i][j+1] = b.y; - vals[i][j+2] = b.z; - vals[i][j+3] = b.w; - } - } - } - - int si_start = 4*(16 * threadIdx.y + threadIdx.x); - for (int so = 0; so < C; so += 32) { - __syncthreads(); - int xmod8 = threadIdx.x % 8; - int xby8 = threadIdx.x / 8; - int xo = 4 * xmod8; - for(int y = 2 * threadIdx.y + xby8; y < 128; y += 32) { - st_vec(&lhs_s[y][xo], ld_vec(inp + y * C + so + xo)); - st_vec(&rhs_s[y][xo], ld_vec(weight + y * C + so + xo)); - } - __syncthreads(); - - for (int si = si_start; si < si_start + 32; si += 4) { - float4 rhs[8]; - for (int u = 0; u < 8; ++u) { - rhs[u] = ld_vec(&rhs_s[u + 8 * threadIdx.y][si % 32]); - } - - for (int ii = 0; ii < 8; ++ii) { - float4 lhs = ld_vec(&lhs_s[ii + 8 * threadIdx.x][si % 32]); - for (int ji = 0; ji < 8; ++ji) { - vals[ii][ji] += lhs.x * rhs[ji].x; - vals[ii][ji] += lhs.y * rhs[ji].y; - vals[ii][ji] += lhs.z * rhs[ji].z; - vals[ii][ji] += lhs.w * rhs[ji].w; - } - } - } - } - - for (int i = 0; i < 8; ++i) { - for (int j = 0; j < 8; j += 4) { - float4 result; - result.x = vals[i][j + 0]; - result.y = vals[i][j + 1]; - result.z = vals[i][j + 2]; - result.w = vals[i][j + 3]; - st_vec(out + (8*threadIdx.x+i) * OC + 8*threadIdx.y + j, result); - } - } -} - - -// ---------------------------------------------------------------------------- -// kernel launchers - -void encoder_forward(float* out, - const int* inp, const float* wte, const float* wpe, - int B, int T, int C) { - assert(C % 4 == 0); - const int block_size = 512; - const int N = B * T * C; - const int grid_size = CEIL_DIV(N / 4, block_size); - encoder_forward_kernel3<<>>((float4*) out, inp, (float4*) wte, (float4*) wpe, B, T, C); - cudaCheck(cudaGetLastError()); -} - -void encoder_backward(float* dwte, float* dwpe, - const float* dout, const int* inp, - int B, int T, int C) { - const int N = B * T * C; - const int block_size = 256; - const int grid_size = CEIL_DIV(N, block_size); - encoder_backward_kernel<<>>(dwte, dwpe, dout, inp, B, T, C); - cudaCheck(cudaGetLastError()); -} - -void layernorm_forward(float* out, float* mean, float* rstd, - float* inp, float* weight, float* bias, - int B, int T, int C) { - const int block_size = 512; - const int N = B * T; - const int grid_size = CEIL_DIV(N * 32, block_size); - layernorm_forward_kernel3<<>>(out, mean, rstd, inp, weight, bias, N, C); - cudaCheck(cudaGetLastError()); -} - -// kernel 1 is the most naive matmul kernel -void matmul_forward(float* out, - const float* inp, const float* weight, const float* bias, - int B, int T, int C, int OC) { - // out is (B,T,OC). OC is short for "output channels", e.g. OC = 4 * C - // inp is (B,T,C), weight is (OC, C), bias is (OC) - int sqrt_block_size = 16; - - dim3 gridDim(CEIL_DIV(B * T, 8*sqrt_block_size), CEIL_DIV(OC, 8*sqrt_block_size)); - dim3 blockDim(sqrt_block_size, sqrt_block_size); - matmul_forward_kernel4<<>>(out, inp, weight, bias, C, OC); - cudaCheck(cudaGetLastError()); -} - -void attention_forward(float* out, float* qkvr, float* att, - float* inp, - int B, int T, int C, int NH) { - // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer. - // Its contents will be overwritten by this function. - const int block_size = 256; - const int softmax_block_size = 256; - - // inp is (B, T, 3C) QKV - // preatt, att are (B, NH, T, T) - // output is (B, T, C) - int HS = C / NH; // head size - - // permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS) - float *q, *k, *v; - q = qkvr + 0 * B * T * C; - k = qkvr + 1 * B * T * C; - v = qkvr + 2 * B * T * C; - int total_threads = B * NH * T * HS; - int num_blocks = CEIL_DIV(total_threads, block_size); - permute_kernel<<>>(q, k, v, inp, B, T, NH, HS); - cudaCheck(cudaGetLastError()); - - // batched matrix multiply with cuBLAS - const float alpha = 1.0f; - const float beta = 0.0f; - float* preatt = inp; - cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha, k, HS, T * HS, q, HS, T * HS, &beta, preatt, T, T * T, B * NH)); - - // 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<<>>(att, scale, preatt, B * NH, T); - cudaCheck(cudaGetLastError()); - - // new approach: first cuBLAS another batched matmul - float* vaccum = inp; - // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs) - cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha, v, HS, T * HS, att, T, T * T, &beta, vaccum, HS, T * HS, B * NH)); - - // 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<<>>(vaccum, out, B, T, NH, HS); - cudaCheck(cudaGetLastError()); -} - -void residual_forward(float* out, float* inp1, float* inp2, int N) { - const int block_size = 256; - const int grid_size = CEIL_DIV(N, block_size); - residual_forward_kernel<<>>(out, inp1, inp2, N); - cudaCheck(cudaGetLastError()); -} - -void gelu_forward(float* out, const float* inp, int N) { - const int block_size = 128; - const int grid_size = CEIL_DIV(N, block_size); - gelu_forward_kernel<<>>(out, inp, N); - cudaCheck(cudaGetLastError()); -} - -void gelu_backward(float* dinp, const float* inp, const float* dout, const int N) { - const int block_size = 128; - const int grid_size = CEIL_DIV(N, block_size); - gelu_backward_kernel<<>>(dinp, inp, dout, N); - cudaCheck(cudaGetLastError()); -} - -void matmul_backward(float* dinp, float* dweight, float* dbias, - float* dout, float* inp, float* weight, - int B, int T, int C, int OC) { - float one = 1.0f; - float zero = 0.0f; - // backward to input, uses = in the backward pass (set the gradient) - cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one, weight, C, dout, OC, &zero, dinp, C)); - // backward to weight, uses += in the backward pass (accumulate the gradient) - cublasCheck(cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one, inp, C, dout, OC, &one, dweight, C)); - // backward to bias, if given, does a += - if (dbias != NULL) { - const int block_size = 1024; - const int grid_size = OC / 32; // for now, OC must be divisible by 32 for this kernel to work - matmul_backward_bias_kernel4<<>>(dbias, dout, B, T, OC); - cudaCheck(cudaGetLastError()); - } -} - -void layernorm_backward(float* dinp, float* dweight, float* dbias, - const float* dout, const float* inp, const float* weight, const float* mean, const float* rstd, - int B, int T, int C) { - const int block_size = 512; - const int N = B * T; - const int grid_size = CEIL_DIV(32*N, block_size); - size_t shared_mem_size = 2 * C * sizeof(float); - layernorm_backward_kernel2<<>>(dinp, dweight, dbias, dout, inp, weight, mean, rstd, B, T, C); - cudaCheck(cudaGetLastError()); -} - -// the sequence of transformations in this compound op is: -// inp (B,T,3C) -> qkvr (B,T,3C) -> preatt (B,NH,T,T) -> att (B,NH,T,T) -> vaccum (B,T,C) -> out (B,T,C) -void attention_backward(float* dinp, float* dqkvr, float* dpreatt, float* datt, float* scratch, - const float* dout, - const float* qkvr, const float* att, - int B, int T, int C, int NH) { - const int block_size = 256; - int HS = C / NH; // head size - const float one = 1.0f; - const float zero = 0.0f; // note beta = 1.0f so that we accumulate gradients (+=) - // unpack convenience pointers into q, k, v - const float *q, *k, *v; - q = qkvr + 0 * B * T * C; - k = qkvr + 1 * B * T * C; - v = qkvr + 2 * B * T * C; - float *dq, *dk, *dv; - dq = dqkvr + 0 * B * T * C; - dk = dqkvr + 1 * B * T * C; - dv = dqkvr + 2 * B * T * C; - // backward through the unpermute operation - int num_blocks = CEIL_DIV(B * T * C, block_size); - unpermute_kernel_backward<<>>(scratch, dout, B, T, NH, HS); - cudaCheck(cudaGetLastError()); - // backward into datt - cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &one, v, HS, T * HS, scratch, HS, T * HS, &zero, datt, T, T * T, B * NH)); - // backward into dv - cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, scratch, HS, T * HS, att, T, T * T, &zero, dv, HS, T * HS, B * NH)); - // backward into preatt - int hs = C / NH; // head size - float scale = 1.0f / sqrtf(hs); - softmax_autoregressive_backward_kernel<<>>(dpreatt, datt, att, B, T, C, scale); - cudaCheck(cudaGetLastError()); - // backward into q - cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &one, k, HS, T * HS, dpreatt, T, T * T, &zero, dq, HS, T * HS, B * NH)); - // backward into k - cublasCheck(cublasSgemmStridedBatched(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &one, q, HS, T * HS, dpreatt, T, T * T, &zero, dk, HS, T * HS, B * NH)); - // backward into inp - num_blocks = CEIL_DIV(B * NH * T * HS, block_size); - permute_kernel_backward<<>>(dinp, dq, dk, dv, B, T, NH, HS); - cudaCheck(cudaGetLastError()); -} - -// replaces logits with logit gradients -void fused_classifier3(float* logits, float* losses, - const float* dlosses, const int* targets, - int B, int T, int V, int P) { - const int block_size = 1024; - const int N = B * T; - const int grid_size = N; - fused_classifier_kernel3<<>>(logits, losses, NULL, dlosses, targets, B, T, V, P); - cudaCheck(cudaGetLastError()); -} - -// ---------------------------------------------------------------------------- -// GPT-2 model definition - -typedef struct { - int max_seq_len; // max sequence length, e.g. 1024 - int vocab_size; // vocab size, e.g. 50257 - int padded_vocab_size; // padded to e.g. %128==0, 50304 - int num_layers; // number of layers, e.g. 12 - int num_heads; // number of heads in attention, e.g. 12 - int channels; // number of channels, e.g. 768 -} GPT2Config; - -// the parameters of the model -#define NUM_PARAMETER_TENSORS 16 -typedef struct { - float* wte; // (V, C) - float* wpe; // (maxT, C) - float* ln1w; // (L, C) - float* ln1b; // (L, C) - float* qkvw; // (L, 3*C, C) - float* qkvb; // (L, 3*C) - float* attprojw; // (L, C, C) - float* attprojb; // (L, C) - float* ln2w; // (L, C) - float* ln2b; // (L, C) - float* fcw; // (L, 4*C, C) - float* fcb; // (L, 4*C) - float* fcprojw; // (L, C, 4*C) - float* fcprojb; // (L, C) - float* lnfw; // (C) - float* lnfb; // (C) -} ParameterTensors; - -void fill_in_parameter_sizes(size_t* param_sizes, GPT2Config config) { - int Vp = config.padded_vocab_size; - int C = config.channels; - int maxT = config.max_seq_len; - int L = config.num_layers; - param_sizes[0] = Vp * C; // wte - param_sizes[1] = maxT * C; // wpe - param_sizes[2] = L * C; // ln1w - param_sizes[3] = L * C; // ln1b - param_sizes[4] = L * (3 * C) * C; // qkvw - param_sizes[5] = L * (3 * C); // qkvb - param_sizes[6] = L * C * C; // attprojw - param_sizes[7] = L * C; // attprojb - param_sizes[8] = L * C; // ln2w - param_sizes[9] = L * C; // ln2b - param_sizes[10] = L * (4 * C) * C; // fcw - param_sizes[11] = L * (4 * C); // fcb - param_sizes[12] = L * C * (4 * C); // fcprojw - param_sizes[13] = L * C; // fcprojb - param_sizes[14] = C; // lnfw - param_sizes[15] = C; // lnfb -} - -// allocate memory for the parameters and point the individual tensors to the right places -float* malloc_and_point_parameters(ParameterTensors* params, size_t* param_sizes, int on_device) { - // on_device: 0 = CPU, 1 = GPU - // calculate the number of parameters - size_t num_parameters = 0; - for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { - num_parameters += param_sizes[i]; - } - // malloc all parameters all at once on the device - float* params_memory; - if (on_device) { - cudaCheck(cudaMalloc((void**)¶ms_memory, num_parameters * sizeof(float))); - } else { - params_memory = (float*)mallocCheck(num_parameters * sizeof(float)); - } - // assign all the tensors their place in the array - float** ptrs[] = { - ¶ms->wte, ¶ms->wpe, ¶ms->ln1w, ¶ms->ln1b, ¶ms->qkvw, ¶ms->qkvb, - ¶ms->attprojw, ¶ms->attprojb, ¶ms->ln2w, ¶ms->ln2b, ¶ms->fcw, ¶ms->fcb, - ¶ms->fcprojw, ¶ms->fcprojb, ¶ms->lnfw, ¶ms->lnfb - }; - float* params_memory_iterator = params_memory; - for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { - *(ptrs[i]) = params_memory_iterator; - params_memory_iterator += param_sizes[i]; - } - return params_memory; -} - -#define NUM_ACTIVATION_TENSORS 21 -typedef struct { - float* encoded; // (B, T, C) - float* ln1; // (L, B, T, C) - float* ln1_mean; // (L, B, T) - float* ln1_rstd; // (L, B, T) - float* atty; // (L, B, T, C) - float* att; // (L, B, NH, T, T) - float* attproj; // (L, B, T, C) - float* residual2; // (L, B, T, C) - float* ln2; // (L, B, T, C) - float* ln2_mean; // (L, B, T) - float* ln2_rstd; // (L, B, T) - float* fch; // (L, B, T, 4*C) - float* fch_gelu; // (L, B, T, 4*C) - float* fcproj; // (L, B, T, C) - float* residual3; // (L, B, T, C) - float* lnf; // (B, T, C) - float* lnf_mean; // (B, T) - float* lnf_rstd; // (B, T) - - float* losses; // (B, T) - // adding these two compared to the CPU .c code, needed for attention kernel as buffers - float* qkvr; // (L, B, T, 3*C) - // in inference mode, this buffer will store the logits - // in training mode, this buffer will contain the *gradients* of the logits. - // during the processing of transformer blocks, we will also use this as a - // general scratchpad buffer. Allocation is made large enough to hold (B, T, 3C), - // (B, NH, T, T), and (B, T, V) shaped tensors. - float* output; -} ActivationTensors; - -void fill_in_activation_sizes(size_t* act_sizes, int B, int T, GPT2Config config) { - size_t Vp = config.padded_vocab_size; - size_t L = config.num_layers; - size_t NH = config.num_heads; - size_t C = config.channels; - act_sizes[0] = B * T * C; // encoded - act_sizes[1] = L * B * T * C; // ln1 - act_sizes[2] = L * B * T; // ln1_mean - act_sizes[3] = L * B * T; // ln1_rstd - act_sizes[4] = L * B * T * C; // atty - act_sizes[5] = L * B * NH * T * T; // att - act_sizes[6] = L * B * T * C; // attproj - act_sizes[7] = L * B * T * C; // residual2 - act_sizes[8] = L * B * T * C; // ln2 - act_sizes[9] = L * B * T; // ln2_mean - act_sizes[10] = L * B * T; // ln2_rstd - act_sizes[11] = L * B * T * 4*C; // fch - act_sizes[12] = L * B * T * 4*C; // fch_gelu - act_sizes[13] = L * B * T * C; // fcproj - act_sizes[14] = L * B * T * C; // residual3 - act_sizes[15] = B * T * C; // lnf - act_sizes[16] = B * T; // lnf_mean - act_sizes[17] = B * T; // lnf_rstd - act_sizes[18] = B * T; // losses - act_sizes[19] = L * B * T * 3*C; // qkvr - act_sizes[20] = B * T * max(3*C, max(NH*T, Vp)); // output / scratch -} - -// Backward pass is conceptually quite different from forward, because we can discard -// the activations of a layer as soon as we're done with it. This lets us aggressively -// reuse memory, so that we need far fewer tensors for backward state. -#define NUM_BACKWARD_TENSORS 3 -typedef struct { - float* bt4c; // (B, T, 4*C) - float* preatt; // (B, NH, T, T) - float* residual3; // (B, T, C) -} GradActTensors; - - -void fill_in_grad_act_sizes(size_t* act_sizes, int B, int T, GPT2Config config) { - size_t NH = config.num_heads; - size_t C = config.channels; - act_sizes[0] = B * T * 4 * C; // bt4c - act_sizes[1] = B * NH * T * T; // preatt - act_sizes[2] = B * T * C; // residual3 -} - - -float* malloc_and_point(float** targets[], const size_t* act_sizes, int n) { - size_t num_activations = 0; - for (size_t i = 0; i < n; i++) { - num_activations += act_sizes[i]; - } - float* acts_memory; - cudaCheck(cudaMalloc((void**)&acts_memory, num_activations * sizeof(float))); - float* acts_memory_iterator = acts_memory; - for (size_t i = 0; i < n; i++) { - *(targets[i]) = acts_memory_iterator; - acts_memory_iterator += act_sizes[i]; - } - return acts_memory; -} - -float* malloc_and_point_activations(ActivationTensors* acts, const size_t* act_sizes) { - float** ptrs[] = { - &acts->encoded, &acts->ln1, &acts->ln1_mean, &acts->ln1_rstd, &acts->atty, - &acts->att, &acts->attproj, &acts->residual2, &acts->ln2, &acts->ln2_mean, - &acts->ln2_rstd, &acts->fch, &acts->fch_gelu, &acts->fcproj, &acts->residual3, &acts->lnf, - &acts->lnf_mean, &acts->lnf_rstd, &acts->losses, &acts->qkvr, &acts->output - }; - return malloc_and_point(ptrs, act_sizes, NUM_ACTIVATION_TENSORS); -} - -float* malloc_and_point_backward(GradActTensors* acts, const size_t* act_sizes) { - float** ptrs[] = { - &acts->bt4c, &acts->preatt, &acts->residual3 - }; - return malloc_and_point(ptrs, act_sizes, NUM_BACKWARD_TENSORS); -} - -typedef struct { - GPT2Config config; - // the weights of the model, and their sizes - ParameterTensors params; - size_t param_sizes[NUM_PARAMETER_TENSORS]; - float* params_memory; - size_t num_parameters; - // gradients of the weights - ParameterTensors grads; - float* grads_memory; - // buffers for the AdamW optimizer - float* m_memory; - float* v_memory; - // the activations of the model, and their sizes - ActivationTensors acts; - size_t act_sizes[NUM_ACTIVATION_TENSORS]; - float* acts_memory; - size_t num_activations; - // gradients of the activations - GradActTensors grads_acts; - size_t num_grad_acts; - float* grads_acts_memory; - // other run state configuration - int batch_size; // the batch size (B) of current forward pass - int seq_len; // the sequence length (T) of current forward pass - int* inputs; // the input tokens for the current forward pass - int* targets; // the target tokens for the current forward pass - float mean_loss; // after a forward pass with targets, will be populated with the mean loss - float* cpu_losses; // CPU buffer to copy the losses to, allocated with cudaMallocHost -} GPT2; - -void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { - - // read in model from a checkpoint file - FILE *model_file = fopenCheck(checkpoint_path, "rb"); - int model_header[256]; - freadCheck(model_header, sizeof(int), 256, model_file); - if (model_header[0] != 20240326) { fprintf(stderr, "Bad magic model file\n"); exit(EXIT_FAILURE); } - if (model_header[1] != 3) { - // was bumped from 1 -> 3 to incorporate the padded vocab size - fprintf(stderr, "Bad version in model file\n"); - fprintf(stderr, "---> HINT: try to re-run `python train_gpt2.py`\n"); - exit(EXIT_FAILURE); - } - - // read in hyperparameters - model->config.max_seq_len = model_header[2]; - model->config.vocab_size = model_header[3]; - model->config.num_layers = model_header[4]; - model->config.num_heads = model_header[5]; - model->config.channels = model_header[6]; - model->config.padded_vocab_size = model_header[7]; - - // allocate space for all the parameters and read them in - fill_in_parameter_sizes(model->param_sizes, model->config); - - // count the number of parameters - size_t num_parameters = 0; - for (size_t i = 0; i < NUM_PARAMETER_TENSORS; i++) { - num_parameters += model->param_sizes[i]; - } - model->num_parameters = num_parameters; - - // create memory for model parameters on the device - model->params_memory = malloc_and_point_parameters(&model->params, model->param_sizes, 1); - - // read in all the parameters from file and copy them to device - float* params_memory_cpu = (float*)mallocCheck(num_parameters * sizeof(float)); - freadCheck(params_memory_cpu, sizeof(float), num_parameters, model_file); - cudaCheck(cudaMemcpy(model->params_memory, params_memory_cpu, num_parameters * sizeof(float), cudaMemcpyHostToDevice)); - free(params_memory_cpu); - fcloseCheck(model_file); - - // other inits - model->acts_memory = NULL; - model->grads_memory = NULL; - model->m_memory = NULL; - model->v_memory = NULL; - model->grads_acts_memory = NULL; - model->inputs = NULL; - model->targets = NULL; - model->cpu_losses = NULL; - model->batch_size = 0; - model->seq_len = 0; - model->mean_loss = -1.0f; // -1.0f will designate no loss -} - -void gpt2_forward(GPT2 *model, int* inputs, int* targets, int B, int T) { - // targets are optional and could be NULL - - // ensure the model was initialized or error out - if (model->params_memory == NULL) { - printf("Error: model was not initialized properly.\n"); - exit(EXIT_FAILURE); - } - - // convenience parameters - int V = model->config.vocab_size; - int Vp = model->config.padded_vocab_size; - int L = model->config.num_layers; - int NH = model->config.num_heads; - int C = model->config.channels; - - // validate inputs, all indices must be in the range [0, V) - for(int i = 0; i < B * T; i++) { - assert(0 <= inputs[i] && inputs[i] < V); - if (targets != NULL) { - assert(0 <= targets[i] && targets[i] < V); - } - } - - // allocate space for all the activations if needed (done here, lazily) - if(model->acts_memory == NULL) { - // record the current B,T as well - model->batch_size = B; - model->seq_len = T; - // and now allocate the space - fill_in_activation_sizes(model->act_sizes, B, T, model->config); - size_t num_activations = 0; - for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) { - num_activations += model->act_sizes[i]; - } - model->num_activations = num_activations; - model->acts_memory = malloc_and_point_activations(&model->acts, model->act_sizes); - printf("allocated %zu MiB for activations\n", (num_activations * sizeof(float)) >> 20); // >> 20 is /(1024*1024) - // also create memory for caching inputs and targets - cudaCheck(cudaMalloc((void**)&model->inputs, B * T * sizeof(int))); - cudaCheck(cudaMalloc((void**)&model->targets, B * T * sizeof(int))); - cudaCheck(cudaMallocHost((void**)&model->cpu_losses, B * T * sizeof(float))); - } else { - // validate B,T is consistent with how we've allocated the memory before - // in principle we could get more clever here in the future, for now this is safest - if (B != model->batch_size || T != model->seq_len) { - printf("Model: B=%d T=%d, Desired: B=%d T=%d\n", model->batch_size, model->seq_len, B, T); - exit(EXIT_FAILURE); - } - } - - // copy inputs/targets to the model - cudaCheck(cudaMemcpy(model->inputs, inputs, B * T * sizeof(int), cudaMemcpyHostToDevice)); - if (targets != NULL) { - cudaCheck(cudaMemcpy(model->targets, targets, B * T * sizeof(int), cudaMemcpyHostToDevice)); - } - - // forward pass - ParameterTensors params = model->params; // for brevity - ActivationTensors acts = model->acts; - float* residual; - encoder_forward(acts.encoded, model->inputs, params.wte, params.wpe, B, T, C); // encoding goes into residual[0] - - for (int l = 0; l < L; l++) { - - residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; - - // get the pointers of the weights for this layer - float* l_ln1w = params.ln1w + l * C; - float* l_ln1b = params.ln1b + l * C; - float* l_qkvw = params.qkvw + l * 3*C * C; - float* l_qkvb = params.qkvb + l * 3*C; - float* l_attprojw = params.attprojw + l * C * C; - float* l_attprojb = params.attprojb + l * C; - float* l_ln2w = params.ln2w + l * C; - float* l_ln2b = params.ln2b + l * C; - float* l_fcw = params.fcw + l * 4*C * C; - float* l_fcb = params.fcb + l * 4*C; - float* l_fcprojw = params.fcprojw + l * C * 4*C; - float* l_fcprojb = params.fcprojb + l * C; - - // get the pointers of the activations for this layer - float* l_ln1 = acts.ln1 + l * B * T * C; - float* l_ln1_mean = acts.ln1_mean + l * B * T; - float* l_ln1_rstd = acts.ln1_rstd + l * B * T; - float* l_qkvr = acts.qkvr + l * B * T * 3*C; - float* l_atty = acts.atty + l * B * T * C; - float* l_att = acts.att + l * B * NH * T * T; - float* l_attproj = acts.attproj + l * B * T * C; - float* l_residual2 = acts.residual2 + l * B * T * C; - float* l_ln2 = acts.ln2 + l * B * T * C; - float* l_ln2_mean = acts.ln2_mean + l * B * T; - float* l_ln2_rstd = acts.ln2_rstd + l * B * T; - float* l_fch = acts.fch + l * B * T * 4*C; - float* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C; - float* l_fcproj = acts.fcproj + l * B * T * C; - float* l_residual3 = acts.residual3 + l * B * T * C; - // these are only needed as scratchpads for the forward pass, but - // need not be stored for backward - float* scratch = acts.output; - - // now do the forward pass - layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C); - matmul_forward(scratch, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C); - attention_forward(l_atty, l_qkvr, l_att, scratch, B, T, C, NH); - matmul_forward(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C); - residual_forward(l_residual2, residual, l_attproj, B*T*C); - layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C); - matmul_forward(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C); - gelu_forward(l_fch_gelu, l_fch, B*T*4*C); - matmul_forward(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C); - residual_forward(l_residual3, l_residual2, l_fcproj, B*T*C); - } - - residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 - layernorm_forward(acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, params.lnfw, params.lnfb, B, T, C); - matmul_forward(acts.output, acts.lnf, params.wte, NULL, B, T, C, Vp); - - // also forward the cross-entropy loss function if we have the targets - if (targets != NULL) { - // 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, 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(float), cudaMemcpyDeviceToHost)); - float mean_loss = 0.0f; - for (int i=0; icpu_losses[i]; } - mean_loss /= B*T; - model->mean_loss = mean_loss; - - } else { - // if we don't have targets, we don't have loss - model->mean_loss = -1.0f; - } -} - -void gpt2_zero_grad(GPT2 *model) { - if (model->grads_acts_memory != NULL) { cudaCheck(cudaMemset(model->grads_acts_memory, 0, model->num_grad_acts * sizeof(float))); } - if (model->grads_memory != NULL) { cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(float))); } -} - -void gpt2_backward(GPT2 *model) { - - // double check we forwarded previously, with targets - if (model->mean_loss == -1.0f) { - printf("Error: must forward with targets before backward\n"); - exit(EXIT_FAILURE); - } - - // lazily allocate the memory for gradients of the weights and activations, if needed - if (model->grads_memory == NULL) { - // allocate buffers for weight gradients - model->grads_memory = malloc_and_point_parameters(&model->grads, model->param_sizes, 1); - printf("allocated %zu MiB for parameter gradients\n", (model->num_parameters * sizeof(float)) >> 20); - // we're going to be clever for the activations backward pass. we don't need to exactly - // mirror the forward pass acrtivations and we will save memory. - size_t bw_act_sizes[NUM_ACTIVATION_TENSORS]; - GPT2Config cfg = model->config; - cfg.num_layers = 1; // copy the configuration but override number of layers to 1 - fill_in_grad_act_sizes(bw_act_sizes, model->batch_size, model->seq_len, cfg); - // count up and allocate the space - model->grads_acts_memory = malloc_and_point_backward(&model->grads_acts, bw_act_sizes); - model->num_grad_acts = 0; - for (int i = 0; i < NUM_BACKWARD_TENSORS; i++) { - model->num_grad_acts += bw_act_sizes[i]; - } - printf("allocated %zu MiB for activation gradients\n", (model->num_grad_acts * sizeof(float)) >> 20); - // init gradients of parameters and activations to zero - gpt2_zero_grad(model); - } - - // convenience shortcuts - int B = model->batch_size; - int T = model->seq_len; - int Vp = model->config.padded_vocab_size; - int L = model->config.num_layers; - int NH = model->config.num_heads; - int C = model->config.channels; - - // backward pass: go in the reverse order of the forward pass, and call backward() functions - ParameterTensors params = model->params; // for brevity - ParameterTensors grads = model->grads; - ActivationTensors acts = model->acts; - GradActTensors grads_acts = model->grads_acts; - - // we kick off the chain rule by filling in dlosses with 1.0f/(B*T) - // this was done in the fused classifier kernel as last step of forward pass - // technically that is a small, inline backward() pass of calculating - // total, final loss as the mean over all losses over all (B,T) positions in the batch - // next: backward the classifier matmul - matmul_backward(grads_acts.bt4c, grads.wte, NULL, acts.output, acts.lnf, params.wte, B, T, C, Vp); - // backward the final layernorm - float* residual = acts.residual3 + (L-1) * B * T * C; // last residual is in residual3 - float* dresidual = grads_acts.residual3; // the main buffer holding the gradient in the backward pass - layernorm_backward(dresidual, grads.lnfw, grads.lnfb, grads_acts.bt4c, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C); - - // now backward all the layers - for (int l = L-1; l >= 0; l--) { - residual = l == 0 ? acts.encoded : acts.residual3 + (l-1) * B * T * C; - - // get the pointers of the weights for this layer - float* l_ln1w = params.ln1w + l * C; - float* l_qkvw = params.qkvw + l * 3*C * C; - float* l_attprojw = params.attprojw + l * C * C; - float* l_ln2w = params.ln2w + l * C; - float* l_fcw = params.fcw + l * 4*C * C; - float* l_fcprojw = params.fcprojw + l * C * 4*C; - // get the pointers of the gradients of the weights for this layer - float* dl_ln1w = grads.ln1w + l * C; - float* dl_ln1b = grads.ln1b + l * C; - float* dl_qkvw = grads.qkvw + l * 3*C * C; - float* dl_qkvb = grads.qkvb + l * 3*C; - float* dl_attprojw = grads.attprojw + l * C * C; - float* dl_attprojb = grads.attprojb + l * C; - float* dl_ln2w = grads.ln2w + l * C; - float* dl_ln2b = grads.ln2b + l * C; - float* dl_fcw = grads.fcw + l * 4*C * C; - float* dl_fcb = grads.fcb + l * 4*C; - float* dl_fcprojw = grads.fcprojw + l * C * 4*C; - float* dl_fcprojb = grads.fcprojb + l * C; - // get the pointers of the activations for this layer - float* l_ln1 = acts.ln1 + l * B * T * C; - float* l_ln1_mean = acts.ln1_mean + l * B * T; - float* l_ln1_rstd = acts.ln1_rstd + l * B * T; - float* l_qkvr = acts.qkvr + l * B * T * 3*C; - float* l_atty = acts.atty + l * B * T * C; - float* l_att = acts.att + l * B * NH * T * T; - float* l_residual2 = acts.residual2 + l * B * T * C; - float* l_ln2 = acts.ln2 + l * B * T * C; - float* l_ln2_mean = acts.ln2_mean + l * B * T; - float* l_ln2_rstd = acts.ln2_rstd + l * B * T; - float* l_fch = acts.fch + l * B * T * 4*C; - float* l_fch_gelu = acts.fch_gelu + l * B * T * 4*C; - // get the pointers of the gradients of the activations for this layer - // notice that there is no l *, because we just have a single copy, and keep - // re-using this memory in every Transformer block as we calculate backward pass - - // we need a B x T x C buffer; thankfully, the forward activation for lnf isn't needed anymore, - // so we can co-opt it here. - float* dl_btc = acts.lnf; - float* dl_bt4c = grads_acts.bt4c; - float* dl_preatt = grads_acts.preatt; - - // re-use scratch buffer of the forward pass - float* scratch = acts.output; - - // backprop this layer - matmul_backward(dl_bt4c, dl_fcprojw, dl_fcprojb, dresidual, l_fch_gelu, l_fcprojw, B, T, 4*C, C); - gelu_backward(dl_bt4c, l_fch, dl_bt4c, B*T*4*C); - matmul_backward(dl_btc, dl_fcw, dl_fcb, dl_bt4c, l_ln2, l_fcw, B, T, C, 4 * C); - // layernorm backward does += to the dresidual, so it correctly accumulates grad from the MLP block above - layernorm_backward(dresidual, dl_ln2w, dl_ln2b, dl_btc, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C); - matmul_backward(dl_btc, dl_attprojw, dl_attprojb, dresidual, l_atty, l_attprojw, B, T, C, C); - // we more B x T x (4)C buffers. l_atty and l_fch aren't needed anymore at this point, so reuse their memory - float* buffer_a = l_atty; - float* buffer_b = l_fch; // this is B x T x 4C, so even larger than what we need - - attention_backward(dl_bt4c, buffer_b, dl_preatt, scratch, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH); - matmul_backward(dl_btc, dl_qkvw, dl_qkvb, dl_bt4c, l_ln1, l_qkvw, B, T, C, 3 * C); - // layernorm backward does += to dresidual, so it correctly accumulates gradient for the Attention block above - layernorm_backward(dresidual, dl_ln1w, dl_ln1b, 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); -} - -void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) { - // reference: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html - - // lazily allocate the memory for m_memory and v_memory - if (model->m_memory == NULL) { - cudaCheck(cudaMalloc((void**)&model->m_memory, model->num_parameters * sizeof(float))); - cudaCheck(cudaMalloc((void**)&model->v_memory, model->num_parameters * sizeof(float))); - cudaCheck(cudaMemset(model->m_memory, 0, model->num_parameters * sizeof(float))); - cudaCheck(cudaMemset(model->v_memory, 0, model->num_parameters * sizeof(float))); - printf("allocated %zu MiB for AdamW optimizer state m\n", (model->num_parameters * sizeof(float)) >> 20); - printf("allocated %zu MiB for AdamW optimizer state v\n", (model->num_parameters * sizeof(float)) >> 20); - } - - int block_size = 512; - int num_blocks = CEIL_DIV(model->num_parameters, block_size); - float beta1_correction = 1.0f - powf(beta1, t); - float beta2_correction = 1.0f - powf(beta2, t); - adamw_kernel2<<>>(model->params_memory, model->grads_memory, model->m_memory, model->v_memory, - model->num_parameters, - learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay); - cudaCheck(cudaGetLastError()); -} - -void gpt2_free(GPT2 *model) { - cudaCheck(cudaFree(model->params_memory)); - cudaCheck(cudaFree(model->grads_memory)); - cudaCheck(cudaFree(model->m_memory)); - cudaCheck(cudaFree(model->v_memory)); - cudaCheck(cudaFree(model->acts_memory)); - cudaCheck(cudaFree(model->grads_acts_memory)); - cudaCheck(cudaFree(model->inputs)); - cudaCheck(cudaFree(model->targets)); - cudaFreeHost(model->cpu_losses); -} - -#ifndef TESTING -// if we are TESTING (see test_gpt2.cu), we'll skip the int main below -// ---------------------------------------------------------------------------- -// sampler: takes probabilities and samples integers from them - -#define GPT2_EOT 50256 - -unsigned int random_u32(unsigned long long *state) { - // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A - *state ^= *state >> 12; - *state ^= *state << 25; - *state ^= *state >> 27; - return (*state * 0x2545F4914F6CDD1Dull) >> 32; -} -float random_f32(unsigned long long *state) { // random float32 in [0,1) - return (random_u32(state) >> 8) / 16777216.0f; -} - -int sample_softmax(const float* logits, int n, float coin) { - // sample index from logits (converted to probabilities using softmax) - // coin is a random number in [0, 1), usually from random_f32() - double norm = 0; - for (int i = 0; i < n; i++) { - norm += expf(logits[i]); - } - // instead of dividing all exp(logits), we can just multiply coin. - coin *= norm; - float cdf = 0.0f; - for (int i = 0; i < n; i++) { - cdf += expf(logits[i]); - if (coin < cdf) { - return i; - } - } - return n - 1; // in case of rounding errors -} - -// ---------------------------------------------------------------------------- -// Logger lite, will probably grow/change some over time - -typedef struct { - FILE *logfile; - int flush_every; // every how many steps to flush the log -} Logger; - -void logger_init(Logger *logger, const char *filename) { - logger->flush_every = 20; - logger->logfile = NULL; - if (filename != NULL) { logger->logfile = fopenCheck(filename, "w"); } -} - -void logger_log_val(Logger *logger, int step, float val_loss) { - if (logger->logfile != NULL) { - fprintf(logger->logfile, "s:%d tel:%.4f\n", step, val_loss); - } -} - -void logger_log_train(Logger *logger, int step, float train_loss) { - if (logger->logfile != NULL) { - fprintf(logger->logfile, "s:%d trl:%.4f\n", step, train_loss); - if (step % 10 == 0) { fflush(logger->logfile); } - } -} - -void logger_free(Logger *logger) { - if (logger->logfile != NULL) { fclose(logger->logfile); } -} - -// ---------------------------------------------------------------------------- -// CLI, poor man's argparse - -void error_usage() { - fprintf(stderr, "Usage: ./train_gpt2fp32cu [options]\n"); - fprintf(stderr, "Options:\n"); - fprintf(stderr, " -i train data filename pattern (default = dev/data/tinyshakespeare/tiny_shakespeare_train.bin)\n"); - fprintf(stderr, " -j val data filename pattern (default = dev/data/tinyshakespeare/tiny_shakespeare_val.bin)\n"); - fprintf(stderr, " -o output log file (default = NULL)\n"); - fprintf(stderr, " -b batch size B (default = 4)\n"); - fprintf(stderr, " -t sequence length T (default = 1024)\n"); - fprintf(stderr, " -l learning rate (default = 3e-4f)\n"); - fprintf(stderr, " -v val_loss_every, how often we evaluate val loss (default = 20)\n"); - fprintf(stderr, " -m val_max_steps, up to how many val batches to estimate val loss? (default = 20)\n"); - fprintf(stderr, " -s sample_every, how often we inference the model (default = 20)\n"); - fprintf(stderr, " -g genT, how many steps of inference we do (default = 64)\n"); - exit(EXIT_FAILURE); -} - -// ---------------------------------------------------------------------------- -// main training loop -int main(int argc, char *argv[]) { - - // read in the (optional) command line arguments - const char* train_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_train.bin"; - const char* val_data_pattern = "dev/data/tinyshakespeare/tiny_shakespeare_val.bin"; - const char* output_log_file = NULL; - int B = 4; // batch size - int T = 1024; // sequence length max - float learning_rate = 3e-4f; - int val_loss_every = 20; // every how many steps do we eval validation loss? - int val_max_steps = 20; // how many batches max do we eval for validation loss? - int sample_every = 20; // every how many steps to do inference? - int genT = 64; // number of steps of inference we will do - for (int i = 1; i < argc; i+=2) { - if (i + 1 >= argc) { error_usage(); } // must have arg after flag - if (argv[i][0] != '-') { error_usage(); } // must start with dash - if (strlen(argv[i]) != 2) { error_usage(); } // must be -x (one dash, one letter) - // read in the args - if (argv[i][1] == 'i') { train_data_pattern = argv[i+1]; } - else if (argv[i][1] == 'j') { val_data_pattern = argv[i+1]; } - else if (argv[i][1] == 'o') { output_log_file = argv[i+1]; } - else if (argv[i][1] == 'b') { B = atoi(argv[i+1]); } - else if (argv[i][1] == 't') { T = atoi(argv[i+1]); } - else if (argv[i][1] == 'l') { learning_rate = atof(argv[i+1]); } - else if (argv[i][1] == 'v') { val_loss_every = atoi(argv[i+1]); } - else if (argv[i][1] == 'm') { val_max_steps = atoi(argv[i+1]); } - else if (argv[i][1] == 's') { sample_every = atoi(argv[i+1]); } - else if (argv[i][1] == 'g') { genT = atoi(argv[i+1]); } - else { error_usage(); } - } - printf("+-----------------------+----------------------------------------------------+\n"); - printf("| Parameter | Value |\n"); - printf("+-----------------------+----------------------------------------------------+\n"); - printf("| train data pattern | %-50s |\n", train_data_pattern); - printf("| val data pattern | %-50s |\n", val_data_pattern); - printf("| output log file | %-50s |\n", output_log_file == NULL ? "NULL" : output_log_file); - printf("| batch size B | %-50d |\n", B); - printf("| sequence length T | %-50d |\n", T); - printf("| learning rate | %-50f |\n", learning_rate); - printf("| val_loss_every | %-50d |\n", val_loss_every); - printf("| val_max_steps | %-50d |\n", val_max_steps); - printf("| sample_every | %-50d |\n", sample_every); - printf("| genT | %-50d |\n", genT); - printf("+-----------------------+----------------------------------------------------+\n"); - - // set up the device - int deviceIdx = 0; - cudaCheck(cudaSetDevice(deviceIdx)); - cudaDeviceProp deviceProp; - cudaGetDeviceProperties(&deviceProp, deviceIdx); - // setup cuBLAS and cuBLASLt - cublasCheck(cublasCreate(&cublas_handle)); - // TF32 precision is equivalent to torch.set_float32_matmul_precision('high') - int enable_tf32 = deviceProp.major >= 8 ? 1 : 0; - cublas_compute_type = enable_tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F; - cublasMath_t cublas_math_mode = enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH; - cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode)); - printf("| device | %-50s |\n", deviceProp.name); - printf("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); - printf("+-----------------------+----------------------------------------------------+\n"); - - // build the GPT-2 model from a checkpoint - GPT2 model; - gpt2_build_from_checkpoint(&model, "gpt2_124M.bin"); - printf("| max_sequence_length T | %-50d |\n", model.config.max_seq_len); - printf("| vocab_size V | %-50d |\n", model.config.vocab_size); - printf("| padded_vocab_size Vp | %-50d |\n", model.config.padded_vocab_size); - printf("| num_layers L | %-50d |\n", model.config.num_layers); - printf("| num_heads NH | %-50d |\n", model.config.num_heads); - printf("| channels C | %-50d |\n", model.config.channels); - printf("| num_parameters | %-50zu |\n", model.num_parameters); - printf("+-----------------------+----------------------------------------------------+\n"); - - // build DataLoaders for both train and val - DataLoader train_loader, val_loader; - dataloader_init(&train_loader, train_data_pattern, B, T, 0, 1, 1); - dataloader_init(&val_loader, val_data_pattern, B, T, 0, 1, 0); - int train_num_batches = train_loader.num_tokens / (B*T); // let's do 1 epoch by default for now - int val_num_batches = val_loader.num_tokens / (B*T); - if (val_num_batches > val_max_steps) { val_num_batches = val_max_steps; } - printf("| train_num_batches | %-50d |\n", train_num_batches); - printf("| val_num_batches | %-50d |\n", val_num_batches); - printf("+-----------------------+----------------------------------------------------+\n"); - - // print model parameter allocations from gpt2_build_from_checkpoint down here to not mess up our table above - printf("allocated %d MiB for model parameters\n", (int)round(model.num_parameters * sizeof(float) / (1024 * 1024))); - - // set up the Logger - Logger logger; - logger_init(&logger, output_log_file); - - // build the Tokenizer - Tokenizer tokenizer; - tokenizer_init(&tokenizer, "gpt2_tokenizer.bin"); - - // some memory for generating samples from the model - unsigned long long rng_state = 1337; - int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int)); - float* cpu_logits = (float*)mallocCheck(model.config.vocab_size * sizeof(float)); - - // train - struct timespec start, end; - double total_sum_iteration_time_s = 0.0; - for (int step = 0; step <= train_num_batches; step++) { - int last_step = step == train_num_batches; - - // once in a while estimate the validation loss - if (step % val_loss_every == 0 || last_step) { - float val_loss = 0.0f; - dataloader_reset(&val_loader); - for (int i = 0; i < val_num_batches; i++) { - dataloader_next_batch(&val_loader); - gpt2_forward(&model, val_loader.inputs, val_loader.targets, B, T); - val_loss += model.mean_loss; - } - val_loss /= val_num_batches; - printf("val loss %f\n", val_loss); - logger_log_val(&logger, step, val_loss); - } - - // once in a while do model inference to print generated text - if (step > 0 && step % sample_every == 0 || last_step) { - // fill up gen_tokens with the GPT2_EOT, which kicks off the generation - for(int i = 0; i < B * T; ++i) { - gen_tokens[i] = GPT2_EOT; - } - // now sample from the model autoregressively - printf("generating:\n---\n"); - for (int t = 1; t < genT; t++) { - // note that inference is very wasteful here because for each token - // we re-calculate the forward pass for all of (B,T) positions from scratch - // but the inference here is just for sanity checking anyway - // and we can maybe optimize a bit more later, with careful tests - gpt2_forward(&model, gen_tokens, NULL, B, T); - // furthermore, below we're only using b=0 (i.e. the first row) of all B rows - // we're in principle running B "inference streams" in parallel here - // only using position 0 because it's a bit faster (copy less probs from GPU -> CPU) - // get the V-dimensional vector probs[0, t-1, :] - float* logits = model.acts.output + (t - 1) * model.config.padded_vocab_size; - // move probs back to CPU and sample (note we only move the first vocab_size logits, ignoring the padding) - cudaCheck(cudaMemcpy(cpu_logits, logits, model.config.vocab_size * sizeof(float), cudaMemcpyDeviceToHost)); - float coin = random_f32(&rng_state); - int next_token = sample_softmax(cpu_logits, model.config.vocab_size, coin); - gen_tokens[t] = next_token; - // print the generated token, either using the Tokenizer or a fallback - if (tokenizer.init_ok) { - const char* token_str = tokenizer_decode(&tokenizer, next_token); - safe_printf(token_str); - } else { - // fall back to printing the token id - printf("%d ", next_token); - } - fflush(stdout); - } - printf("\n---\n"); - } - - // bit confusing: we want to make sure to eval and sample on 0th iteration - // but also after the very last iteration. so we loop for step <= train_num_batches - // instead of just < train_num_batches (one extra due to <=), only to do - // the validation/sampling one last time, and then we break right here as we're done. - if (last_step) { break; } - - // do a training step - clock_gettime(CLOCK_MONOTONIC, &start); - dataloader_next_batch(&train_loader); - gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T); - gpt2_zero_grad(&model); - gpt2_backward(&model); - gpt2_update(&model, learning_rate, 0.9f, 0.999f, 1e-8f, 0.0f, step+1); - cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings - clock_gettime(CLOCK_MONOTONIC, &end); - double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9; - total_sum_iteration_time_s += time_elapsed_s; - int tokens_per_second = (B * T) / time_elapsed_s; - printf("step %4d/%d: train loss %f (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, time_elapsed_s * 1000, tokens_per_second); - logger_log_train(&logger, step, model.mean_loss); - } - // add a total average, for optimizations that are only mild improvements - printf("total average iteration time: %f ms\n", total_sum_iteration_time_s / train_num_batches * 1000); - - // free - dataloader_free(&train_loader); - dataloader_free(&val_loader); - tokenizer_free(&tokenizer); - gpt2_free(&model); - free(cpu_logits); - free(gen_tokens); - cublasCheck(cublasDestroy(cublas_handle)); - logger_free(&logger); - - return 0; -} -#endif \ No newline at end of file