Showing preview only (1,261K chars total). Download the full file or copy to clipboard to get everything.
Repository: karpathy/llm.c
Branch: master
Commit: f1e2ace65149
Files: 102
Total size: 1.2 MB
Directory structure:
gitextract_pp0afoji/
├── .github/
│ └── workflows/
│ ├── ci.yml
│ ├── ci_gpu.yml
│ └── ci_tests.yml
├── .gitignore
├── LICENSE
├── Makefile
├── README.md
├── dev/
│ ├── cpu/
│ │ └── matmul_forward.c
│ ├── cuda/
│ │ ├── Makefile
│ │ ├── README.md
│ │ ├── adamw.cu
│ │ ├── attention_backward.cu
│ │ ├── attention_forward.cu
│ │ ├── benchmark_on_modal.py
│ │ ├── classifier_fused.cu
│ │ ├── common.h
│ │ ├── crossentropy_forward.cu
│ │ ├── crossentropy_softmax_backward.cu
│ │ ├── encoder_backward.cu
│ │ ├── encoder_forward.cu
│ │ ├── fused_residual_forward.cu
│ │ ├── gelu_backward.cu
│ │ ├── gelu_forward.cu
│ │ ├── global_norm.cu
│ │ ├── layernorm_backward.cu
│ │ ├── layernorm_forward.cu
│ │ ├── matmul_backward.cu
│ │ ├── matmul_backward_bias.cu
│ │ ├── matmul_forward.cu
│ │ ├── nccl_all_reduce.cu
│ │ ├── permute.cu
│ │ ├── residual_forward.cu
│ │ ├── softmax_forward.cu
│ │ └── trimat_forward.cu
│ ├── data/
│ │ ├── README.md
│ │ ├── data_common.py
│ │ ├── edu_fineweb.sh
│ │ ├── fineweb.py
│ │ ├── fineweb.sh
│ │ ├── hellaswag.py
│ │ ├── mmlu.py
│ │ ├── tinyshakespeare.py
│ │ └── tinystories.py
│ ├── download_starter_pack.sh
│ ├── eval/
│ │ ├── README.md
│ │ ├── export_hf.py
│ │ ├── run_eval.sh
│ │ └── summarize_eval.py
│ ├── loss_checker_ci.py
│ ├── test/
│ │ ├── Makefile
│ │ ├── device_file_io.cu
│ │ ├── test_dataloader.c
│ │ └── test_outlier_detector.c
│ ├── unistd.h
│ └── vislog.ipynb
├── doc/
│ └── layernorm/
│ ├── layernorm.c
│ ├── layernorm.md
│ └── layernorm.py
├── llmc/
│ ├── adamw.cuh
│ ├── attention.cuh
│ ├── cublas_common.h
│ ├── cuda_common.h
│ ├── cuda_utils.cuh
│ ├── cudnn_att.cpp
│ ├── cudnn_att.h
│ ├── dataloader.h
│ ├── encoder.cuh
│ ├── fused_classifier.cuh
│ ├── gelu.cuh
│ ├── global_norm.cuh
│ ├── layernorm.cuh
│ ├── logger.h
│ ├── matmul.cuh
│ ├── mfu.h
│ ├── outlier_detector.h
│ ├── rand.h
│ ├── sampler.h
│ ├── schedulers.h
│ ├── tokenizer.h
│ ├── utils.h
│ └── zero.cuh
├── profile_gpt2.cu
├── profile_gpt2cu.py
├── requirements.txt
├── scripts/
│ ├── README.md
│ ├── multi_node/
│ │ ├── run_gpt2_124M_fs.sbatch
│ │ ├── run_gpt2_124M_mpi.sh
│ │ └── run_gpt2_124M_tcp.sbatch
│ ├── pyrun_gpt2_124M.sh
│ ├── run_gpt2_124M.sh
│ ├── run_gpt2_1558M.sh
│ ├── run_gpt2_350M.sh
│ ├── run_gpt2_774M.sh
│ └── run_gpt3_125M.sh
├── test_gpt2.c
├── test_gpt2.cu
├── test_gpt2_fp32.cu
├── train_gpt2.c
├── train_gpt2.cu
├── train_gpt2.py
├── train_gpt2_fp32.cu
└── train_llama3.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .github/workflows/ci.yml
================================================
name: Build and test
on:
create:
workflow_dispatch:
push:
branches:
- master
pull_request:
branches:
- master
jobs:
build-and-test-cpu:
strategy:
matrix:
os: [ubuntu-latest, macos-latest, windows-latest]
runs-on: ${{ matrix.os }}
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install OpenMP
if: matrix.os != 'windows-latest'
run: |
if [ "${{ runner.os }}" == "Linux" ]; then
sudo apt-get update && sudo apt-get install -y libomp-dev
elif [ "${{ runner.os }}" == "macOS" ]; then
brew install libomp
fi
- name: Install dependencies
run: pip install -r requirements.txt
- name: Run preprocessing
run: python dev/data/tinyshakespeare.py
- name: Train model
run: python train_gpt2.py --device=cpu
- name: Download Win32 Make.exe
if: matrix.os == 'windows-latest'
run: |
$wc = New-Object System.Net.WebClient
$url = 'https://github.com/maweil/MakeForWindows/releases/download/v4.4.1/make-bin-win64.zip'
$output = './make-bin-win64.zip'
$wc.DownloadFile($url, $output)
- name: Unzip Win32 Makefile
if: matrix.os == 'windows-latest'
run: |
unzip make-bin-win64.zip
- name: Compile training and testing program
if: matrix.os != 'windows-latest'
run: make test_gpt2 train_gpt2
- name: Compile training and testing program for Windows
if: matrix.os == 'windows-latest'
shell: cmd
run: |
call "C:\\Program Files\\Microsoft Visual Studio\\2022\\Enterprise\\VC\\Auxiliary\\Build\\vcvars64.bat"
make-4.4.1\dist\make WIN_CI_BUILD=1 test_gpt2 train_gpt2
- name: Execute testing program (With OpenMP)
if: matrix.os != 'windows-latest'
run: OMP_NUM_THREADS=8 ./test_gpt2
- name: Execute Windows testing program (With OpenMP)
if: matrix.os == 'windows-latest'
shell: cmd
run: |
copy test_gpt2 test_gpt2.exe
test_gpt2.exe
- name: Compile training and testing program without OpenMP
if: matrix.os != 'windows-latest'
run: NO_OMP=1 make test_gpt2 train_gpt2
- name: Execute testing program (No OpenMP)
if: matrix.os != 'windows-latest'
run: ./test_gpt2
build-cuda-windows:
runs-on: windows-latest
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Download Win32 Make.exe
run: |
$wc = New-Object System.Net.WebClient
$url = 'https://github.com/maweil/MakeForWindows/releases/download/v4.4.1/make-bin-win64.zip'
$output = './make-bin-win64.zip'
$wc.DownloadFile($url, $output)
- name: Unzip Win32 Makefile
run: |
unzip make-bin-win64.zip
- name: Install Cuda Toolkit 12.4 on Windows
run: |
mkdir -p "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4"
choco install unzip -y
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/windows-x86_64/cuda_cudart-windows-x86_64-12.4.127-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvcc/windows-x86_64/cuda_nvcc-windows-x86_64-12.4.131-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/windows-x86_64/cuda_nvrtc-windows-x86_64-12.4.127-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/libcublas/windows-x86_64/libcublas-windows-x86_64-12.4.5.8-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvtx/windows-x86_64/cuda_nvtx-windows-x86_64-12.4.127-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_profiler_api/windows-x86_64/cuda_profiler_api-windows-x86_64-12.4.127-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/visual_studio_integration/windows-x86_64/visual_studio_integration-windows-x86_64-12.4.127-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvprof/windows-x86_64/cuda_nvprof-windows-x86_64-12.4.127-archive.zip"
curl -O "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cccl/windows-x86_64/cuda_cccl-windows-x86_64-12.4.127-archive.zip"
unzip '*.zip' -d "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4"
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\cuda_cudart-windows-x86_64-12.4.127-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\cuda_nvcc-windows-x86_64-12.4.131-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\cuda_nvrtc-windows-x86_64-12.4.127-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\libcublas-windows-x86_64-12.4.5.8-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\cuda_nvtx-windows-x86_64-12.4.127-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\cuda_profiler_api-windows-x86_64-12.4.127-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\visual_studio_integration-windows-x86_64-12.4.127-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\cuda_nvprof-windows-x86_64-12.4.127-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
xcopy "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\cuda_cccl-windows-x86_64-12.4.127-archive\*" "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" /E /I /H /Y
# Default installation path for CUDA Toolkit is C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4
- name: Add Path
run: |
echo "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v12.4\\bin" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
echo "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4\libnvvp" | Out-File -FilePath $env:GITHUB_PATH -Encoding utf8 -Append
echo "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
echo "CUDA_PATH_V12_4=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.4" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
- name: Build Cuda targets
shell: cmd
working-directory: ${{ github.workspace }}
run: |
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:
image: nvidia/cuda:12.4.1-devel-ubuntu22.04
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Build FP32 checkpoint
run: make train_gpt2fp32cu test_gpt2fp32cu
- name: Build FP32 precision
run: PRECISION=FP32 make train_gpt2cu test_gpt2cu profile_gpt2cu
build-cuda-bf16:
runs-on: ubuntu-latest
container:
image: nvidia/cuda:12.4.1-devel-ubuntu22.04
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Build project
run: PRECISION=BF16 make test_gpt2cu train_gpt2cu profile_gpt2cu
build-cuda-fp16:
runs-on: ubuntu-latest
container:
image: nvidia/cuda:12.4.1-devel-ubuntu22.04
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Build project
run: PRECISION=FP16 make test_gpt2cu train_gpt2cu profile_gpt2cu
build-cuda-kernels:
runs-on: ubuntu-latest
container:
image: nvidia/cuda:12.4.1-devel-ubuntu22.04
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install OpenMP and OpenMPI
run: apt-get update && apt-get install -y libomp-dev libopenmpi-dev
- name: Build project
run: make -j4 -C dev/cuda
================================================
FILE: .github/workflows/ci_gpu.yml
================================================
name: GPU Builds and Tests
on:
create:
workflow_dispatch:
push:
branches:
- master
pull_request:
branches:
- master
jobs:
build-and-test-gpu:
runs-on: ubicloud-gpu-standard-1-latest
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install OpenMP
run: sudo apt-get update && sudo apt-get install -y libomp-dev
- name: Install dependencies
run: pip install -r requirements.txt
- name: Run preprocessing
run: python dev/data/tinyshakespeare.py
- name: Train model
run: python train_gpt2.py
- name: Compile training and testing program
run: make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu
- name: Train model (With OpenMP)
run: OMP_NUM_THREADS=8 ./train_gpt2cu
- name: Train model (FP32) with gpt2_124M.bin
run: |
PRECISION=FP32 make train_gpt2cu
./train_gpt2cu -b 1 -t 64 -d 256 -l 0.0001 -v 200 -s 200 -a 1 -x 10 -r 0 -f 0 -e "gpt2_124M.bin"
- name: Test for percent loss differential for FP32
run: |
PRECISION=FP32 make train_gpt2cu
./train_gpt2cu -b 1 -t 64 -d 256 -l 0.0001 -v 200 -s 200 -a 1 -x 10 -r 0 -f 0 -e "gpt2_124M.bin" > train_gpt2cu_fp32_precision.txt
python dev/loss_checker_ci.py -f train_gpt2cu_fp32_precision.txt -s 20 -e 28 -a 5.0
- name: Build FP32 precision
run: PRECISION=FP32 make test_gpt2cu profile_gpt2cu
- name: Run default
run: ./test_gpt2cu
- name: Run no recompute GeLU
run: ./test_gpt2cu -r 0
- name: Run recompute LN
run: ./test_gpt2cu -r 2
- name: Build BF16 precision
run: PRECISION=BF16 make train_gpt2cu test_gpt2cu profile_gpt2cu
- name: Run default
run: ./test_gpt2cu
- name: Run no recompute GeLU
run: ./test_gpt2cu -r 0
- name: Run no master weights
run: ./test_gpt2cu -w 0
- name: Run recompute LN
run: ./test_gpt2cu -r 2
- name: Train model fp32 (With OpenMP)
run: OMP_NUM_THREADS=8 ./train_gpt2fp32cu
- name: Execute testing program (With OpenMP)
run: OMP_NUM_THREADS=8 ./test_gpt2cu
- name: Execute testing program fp32 (With OpenMP)
run: OMP_NUM_THREADS=8 ./test_gpt2fp32cu
- name: Compile training and testing program without OpenMP
run: NO_OMP=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu
- name: Train model (No OpenMP)
run: NO_OMP=1 ./train_gpt2cu
- name: Train model fp32 (No OpenMP)
run: NO_OMP=1 ./train_gpt2fp32cu
- name: Execute testing program (No OpenMP)
run: ./test_gpt2cu -b 32
- name: Execute testing program fp32 (No OpenMP)
run: ./test_gpt2fp32cu
- name: Install cuDNN-frontend
run:
git clone https://github.com/NVIDIA/cudnn-frontend.git
- name: Build with cuDNN
run: USE_CUDNN=1 make test_gpt2cu train_gpt2cu test_gpt2fp32cu train_gpt2fp32cu
- name: Train model with cuDNN
run: ./train_gpt2cu
- name: Train model fp32 with cuDNN
run: ./train_gpt2fp32cu
- name: Execute testing program with cuDNN
run: ./test_gpt2cu
- name: Execute testing program fp32 with cuDNN
run: ./test_gpt2fp32cu
unit-tests-gpu:
runs-on: ubicloud-gpu-standard-1-latest
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Test Device<->File IO
run: cd dev/test && nvcc -o device_file_io device_file_io.cu && ./device_file_io
================================================
FILE: .github/workflows/ci_tests.yml
================================================
name: Unit, Static and other Tests
on:
create:
workflow_dispatch:
push:
branches:
- master
pull_request:
branches:
- master
jobs:
dataloader_test:
runs-on: ubuntu-latest
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: test the dataloader without / with sanitize address
run: |
cd dev/test
make PRECISION=BF16 test_dataloader
./test_dataloader
make clean
make PRECISION=BF16 TEST_CFLAGS="-fsanitize=address -fno-omit-frame-pointer" test_dataloader
./test_dataloader
ptx_and_sass_files:
runs-on: ubuntu-latest
container:
image: nvidia/cuda:12.4.1-devel-ubuntu22.04
steps:
- name: Checkout code
uses: actions/checkout@v4
- name: Install OpenMP and OpenMPI
run: apt-get update && apt-get install -y libomp-dev libopenmpi-dev
- name: Generate ptx/sass files and upload them to persistent storage
run: |
mkdir -p dev/cuda/ptx_sass_logs
make train_gpt2cu
cuobjdump --dump-ptx train_gpt2cu > dev/cuda/train_gpt2cu.ptx
cuobjdump --dump-sass train_gpt2cu > dev/cuda/train_gpt2cu.sass
cd dev/cuda
make -j all_ptx
make -j all_sass
cp *.ptx ptx_sass_logs/
cp *.sass ptx_sass_logs/
ls ptx_sass_logs/
- name: Generate ptx/sass files for A100 and upload them to persistent storage
run: |
mkdir -p dev/cuda/ptx_sass_logs_A100
make train_gpt2cu GPU_COMPUTE_CAPABILITY=80
cuobjdump --dump-ptx train_gpt2cu > dev/cuda/train_gpt2cu.ptx
cuobjdump --dump-sass train_gpt2cu > dev/cuda/train_gpt2cu.sass
cd dev/cuda
make -j GPU_COMPUTE_CAPABILITY=80 all_ptx
make -j GPU_COMPUTE_CAPABILITY=80 all_sass
cp *.ptx ptx_sass_logs_A100/
cp *.sass ptx_sass_logs_A100/
ls ptx_sass_logs_A100/
- name: Generate ptx/sass files for H100 and upload them to persistent storage
run: |
mkdir -p dev/cuda/ptx_sass_logs_H100
make train_gpt2cu GPU_COMPUTE_CAPABILITY=90
cuobjdump --dump-ptx train_gpt2cu > dev/cuda/train_gpt2cu.ptx
cuobjdump --dump-sass train_gpt2cu > dev/cuda/train_gpt2cu.sass
cd dev/cuda
make -j GPU_COMPUTE_CAPABILITY=90 all_ptx
make -j GPU_COMPUTE_CAPABILITY=90 all_sass
cp *.ptx ptx_sass_logs_H100/
cp *.sass ptx_sass_logs_H100/
ls ptx_sass_logs_H100/
- name: Upload ptx/sass files
uses: actions/upload-artifact@v4
with:
name: ptx_sass_files
path: dev/cuda/ptx_sass_logs/
retention-days: 30 # days to retain
- name: Upload ptx/sass files for A100
uses: actions/upload-artifact@v4
with:
name: ptx_sass_files_A100
path: dev/cuda/ptx_sass_logs_A100/
retention-days: 30 # days to retain
- name: Upload ptx/sass files for H100
uses: actions/upload-artifact@v4
with:
name: ptx_sass_files_H100
path: dev/cuda/ptx_sass_logs_H100/
retention-days: 30 # days to retain
================================================
FILE: .gitignore
================================================
# dot files and such
.vscode
.venv
# .bin files generated by Python
*.bin
# data directories
dev/data/__pycache__/
dev/data/fineweb10B/
dev/data/hellaswag/
dev/data/mmlu/
dev/data/tinyshakespeare/
dev/data/tinystories/
# binaries
test_gpt2
test_gpt2cu
test_gpt2fp32cu
train_gpt2
train_gpt2cu
train_gpt2fp32cu
profile_gpt2cu
dev/cuda/*_forward
dev/cuda/*_backward
dev/cuda/classifier_fused
dev/cuda/adamw
dev/cuda/matmul_backward_bias
dev/cuda/nccl_all_reduce
dev/cuda/global_norm
*.obj
*.exe
*.o
# log files
*.log
================================================
FILE: LICENSE
================================================
MIT License
Copyright (c) 2024 Andrej Karpathy
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: Makefile
================================================
CC ?= clang
CFLAGS = -Ofast -Wno-unused-result -Wno-ignored-pragmas -Wno-unknown-attributes
LDFLAGS =
LDLIBS = -lm
INCLUDES =
CFLAGS_COND = -march=native
# Find nvcc
SHELL_UNAME = $(shell uname)
REMOVE_FILES = rm -f
OUTPUT_FILE = -o $@
CUDA_OUTPUT_FILE = -o $@
# Default O3 CPU optimization level for NVCC (0 for fastest compile time)
FORCE_NVCC_O ?= 3
# NVCC flags
# -t=0 is short for --threads, 0 = number of CPUs on the machine
NVCC_FLAGS = --threads=0 -t=0 --use_fast_math -std=c++17 -O$(FORCE_NVCC_O)
NVCC_LDFLAGS = -lcublas -lcublasLt
NVCC_INCLUDES =
NVCC_LDLIBS =
NCLL_INCUDES =
NVCC_CUDNN =
# By default we don't build with cudnn because it blows up compile time from a few seconds to ~minute
USE_CUDNN ?= 0
# We will place .o files in the `build` directory (create it if it doesn't exist)
BUILD_DIR = build
ifeq ($(OS), Windows_NT)
$(shell if not exist $(BUILD_DIR) mkdir $(BUILD_DIR))
REMOVE_BUILD_OBJECT_FILES := del $(BUILD_DIR)\*.obj
else
$(shell mkdir -p $(BUILD_DIR))
REMOVE_BUILD_OBJECT_FILES := rm -f $(BUILD_DIR)/*.o
endif
# Function to check if a file exists in the PATH
ifneq ($(OS), Windows_NT)
define file_exists_in_path
$(which $(1) 2>/dev/null)
endef
else
define file_exists_in_path
$(shell where $(1) 2>nul)
endef
endif
ifneq ($(CI),true) # if not in CI, then use the GPU query
ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY=
ifneq ($(call file_exists_in_path, nvidia-smi),)
# Get the compute capabilities of all GPUs
# Remove decimal points, sort numerically in ascending order, and select the first (lowest) value
GPU_COMPUTE_CAPABILITY=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed 's/\.//g' | sort -n | head -n 1)
GPU_COMPUTE_CAPABILITY := $(strip $(GPU_COMPUTE_CAPABILITY))
endif
endif
endif
# set to defaults if - make GPU_COMPUTE_CAPABILITY= otherwise use the compute capability detected above
ifneq ($(GPU_COMPUTE_CAPABILITY),)
NVCC_FLAGS += --generate-code arch=compute_$(GPU_COMPUTE_CAPABILITY),code=[compute_$(GPU_COMPUTE_CAPABILITY),sm_$(GPU_COMPUTE_CAPABILITY)]
endif
# autodect a lot of various supports on current platform
$(info ---------------------------------------------)
ifneq ($(OS), Windows_NT)
NVCC := $(shell which nvcc 2>/dev/null)
NVCC_LDFLAGS += -lnvidia-ml
# Function to test if the compiler accepts a given flag.
define check_and_add_flag
$(eval FLAG_SUPPORTED := $(shell printf "int main() { return 0; }\n" | $(CC) $(1) -x c - -o /dev/null 2>/dev/null && echo 'yes'))
ifeq ($(FLAG_SUPPORTED),yes)
CFLAGS += $(1)
endif
endef
# Check each flag and add it if supported
$(foreach flag,$(CFLAGS_COND),$(eval $(call check_and_add_flag,$(flag))))
else
CFLAGS :=
REMOVE_FILES = del *.exe,*.obj,*.lib,*.exp,*.pdb && del
SHELL_UNAME := Windows
ifneq ($(shell where nvcc 2> nul),"")
NVCC := nvcc
else
NVCC :=
endif
CC := cl
CFLAGS = /Idev /Zi /nologo /W4 /WX- /diagnostics:column /sdl /O2 /Oi /Ot /GL /D _DEBUG /D _CONSOLE /D _UNICODE /D UNICODE /Gm- /EHsc /MD /GS /Gy /fp:fast /Zc:wchar_t /Zc:forScope /Zc:inline /permissive- \
/external:W3 /Gd /TP /wd4996 /Fd$@.pdb /FC /openmp:llvm
LDFLAGS :=
LDLIBS :=
INCLUDES :=
NVCC_FLAGS += -I"dev"
ifeq ($(WIN_CI_BUILD),1)
$(info Windows CI build)
OUTPUT_FILE = /link /OUT:$@
CUDA_OUTPUT_FILE = -o $@
else
$(info Windows local build)
OUTPUT_FILE = /link /OUT:$@ && copy /Y $@ $@.exe
CUDA_OUTPUT_FILE = -o $@ && copy /Y $@.exe $@
endif
endif
# Check and include cudnn if available
# You can override the path to cudnn frontend by setting CUDNN_FRONTEND_PATH on the make command line
# By default, we look for it in HOME/cudnn-frontend/include and ./cudnn-frontend/include
# Refer to the README for cuDNN install instructions
ifeq ($(USE_CUDNN), 1)
ifeq ($(SHELL_UNAME), Linux)
ifeq ($(shell [ -d $(HOME)/cudnn-frontend/include ] && echo "exists"), exists)
$(info ✓ cuDNN found, will run with flash-attention)
CUDNN_FRONTEND_PATH ?= $(HOME)/cudnn-frontend/include
else ifeq ($(shell [ -d cudnn-frontend/include ] && echo "exists"), exists)
$(info ✓ cuDNN found, will run with flash-attention)
CUDNN_FRONTEND_PATH ?= cudnn-frontend/include
else
$(error ✗ cuDNN not found. See the README for install instructions and the Makefile for hard-coded paths)
endif
NVCC_INCLUDES += -I$(CUDNN_FRONTEND_PATH)
NVCC_LDFLAGS += -lcudnn
NVCC_FLAGS += -DENABLE_CUDNN
NVCC_CUDNN = $(BUILD_DIR)/cudnn_att.o
else
ifneq ($(OS), Windows_NT)
$(info → cuDNN is not supported on MAC OS right now)
else
$(info ✓ Windows cuDNN found, will run with flash-attention)
ifeq ($(shell if exist "$(HOMEDRIVE)$(HOMEPATH)\cudnn-frontend\include" (echo exists)),exists)
CUDNN_FRONTEND_PATH ?= $(HOMEDRIVE)$(HOMEPATH)\cudnn-frontend\include #override on command line if different location
else ifeq ($(shell if exist "cudnn-frontend\include" (echo exists)),exists)
CUDNN_FRONTEND_PATH ?= cudnn-frontend\include #override on command line if different location
else
$(error ✗ cuDNN not found. See the README for install instructions and the Makefile for hard-coded paths)
endif
CUDNN_INCLUDE_PATH ?= -I"C:\Program Files\NVIDIA\CUDNN\v9.1\include\12.4"
CUDNN_FRONTEND_PATH += $(CUDNN_INCLUDE_PATH)
NVCC_FLAGS += --std c++20 -Xcompiler "/std:c++20" -Xcompiler "/EHsc /W0 /nologo /Ox /FS" -maxrregcount=0 --machine 64
NVCC_CUDNN = $(BUILD_DIR)\cudnn_att.obj
NVCC_INCLUDES += -I$(CUDNN_FRONTEND_PATH)
NVCC_LDFLAGS += -L"C:\Program Files\NVIDIA\CUDNN\v9.1\lib\12.4\x64" -lcudnn
NVCC_FLAGS += -DENABLE_CUDNN
endif
endif
else
$(info → cuDNN is manually disabled by default, run make with `USE_CUDNN=1` to try to enable)
endif
# Check if OpenMP is available
# This is done by attempting to compile an empty file with OpenMP flags
# OpenMP makes the code a lot faster so I advise installing it
# e.g. on MacOS: brew install libomp
# e.g. on Ubuntu: sudo apt-get install libomp-dev
# later, run the program by prepending the number of threads, e.g.: OMP_NUM_THREADS=8 ./gpt2
# First, check if NO_OMP is set to 1, if not, proceed with the OpenMP checks
ifeq ($(NO_OMP), 1)
$(info OpenMP is manually disabled)
else
ifneq ($(OS), Windows_NT)
# Detect if running on macOS or Linux
ifeq ($(SHELL_UNAME), Darwin)
# Check for Homebrew's libomp installation in different common directories
ifeq ($(shell [ -d /opt/homebrew/opt/libomp/lib ] && echo "exists"), exists)
# macOS with Homebrew on ARM (Apple Silicon)
CFLAGS += -Xclang -fopenmp -DOMP
LDFLAGS += -L/opt/homebrew/opt/libomp/lib
LDLIBS += -lomp
INCLUDES += -I/opt/homebrew/opt/libomp/include
$(info ✓ OpenMP found)
else ifeq ($(shell [ -d /usr/local/opt/libomp/lib ] && echo "exists"), exists)
# macOS with Homebrew on Intel
CFLAGS += -Xclang -fopenmp -DOMP
LDFLAGS += -L/usr/local/opt/libomp/lib
LDLIBS += -lomp
INCLUDES += -I/usr/local/opt/libomp/include
$(info ✓ OpenMP found)
else
$(info ✗ OpenMP not found)
endif
else
# Check for OpenMP support in GCC or Clang on Linux
ifeq ($(shell echo | $(CC) -fopenmp -x c -E - > /dev/null 2>&1; echo $$?), 0)
CFLAGS += -fopenmp -DOMP
LDLIBS += -lgomp
$(info ✓ OpenMP found)
else
$(info ✗ OpenMP not found)
endif
endif
endif
endif
# Check if NCCL is available, include if so, for multi-GPU training
ifeq ($(NO_MULTI_GPU), 1)
$(info → Multi-GPU (NCCL) is manually disabled)
else
ifneq ($(OS), Windows_NT)
# Detect if running on macOS or Linux
ifeq ($(SHELL_UNAME), Darwin)
$(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping NCCL support)
else ifeq ($(shell dpkg -l | grep -q nccl && echo "exists"), exists)
$(info ✓ NCCL found, OK to train with multiple GPUs)
NVCC_FLAGS += -DMULTI_GPU
NVCC_LDLIBS += -lnccl
else
$(info ✗ NCCL is not found, disabling multi-GPU support)
$(info ---> On Linux you can try install NCCL with `sudo apt install libnccl2 libnccl-dev`)
endif
endif
endif
# Attempt to find and include OpenMPI on the system
OPENMPI_DIR ?= /usr/lib/x86_64-linux-gnu/openmpi
OPENMPI_LIB_PATH = $(OPENMPI_DIR)/lib/
OPENMPI_INCLUDE_PATH = $(OPENMPI_DIR)/include/
ifeq ($(NO_USE_MPI), 1)
$(info → MPI is manually disabled)
else ifeq ($(shell [ -d $(OPENMPI_LIB_PATH) ] && [ -d $(OPENMPI_INCLUDE_PATH) ] && echo "exists"), exists)
$(info ✓ MPI enabled)
NVCC_INCLUDES += -I$(OPENMPI_INCLUDE_PATH)
NVCC_LDFLAGS += -L$(OPENMPI_LIB_PATH)
NVCC_LDLIBS += -lmpi
NVCC_FLAGS += -DUSE_MPI
else
$(info ✗ MPI not found)
endif
# Precision settings, default to bf16 but ability to override
PRECISION ?= BF16
VALID_PRECISIONS := FP32 FP16 BF16
ifeq ($(filter $(PRECISION),$(VALID_PRECISIONS)),)
$(error Invalid precision $(PRECISION), valid precisions are $(VALID_PRECISIONS))
endif
ifeq ($(PRECISION), FP32)
PFLAGS = -DENABLE_FP32
else ifeq ($(PRECISION), FP16)
PFLAGS = -DENABLE_FP16
else
PFLAGS = -DENABLE_BF16
endif
# PHONY means these targets will always be executed
.PHONY: all train_gpt2 test_gpt2 train_gpt2cu test_gpt2cu train_gpt2fp32cu test_gpt2fp32cu profile_gpt2cu
# Add targets
TARGETS = train_gpt2 test_gpt2
# Conditional inclusion of CUDA targets
ifeq ($(NVCC),)
$(info ✗ nvcc not found, skipping GPU/CUDA builds)
else
$(info ✓ nvcc found, including GPU/CUDA support)
TARGETS += train_gpt2cu test_gpt2cu train_gpt2fp32cu test_gpt2fp32cu $(NVCC_CUDNN)
endif
$(info ---------------------------------------------)
all: $(TARGETS)
train_gpt2: train_gpt2.c
$(CC) $(CFLAGS) $(INCLUDES) $(LDFLAGS) $^ $(LDLIBS) $(OUTPUT_FILE)
test_gpt2: test_gpt2.c
$(CC) $(CFLAGS) $(INCLUDES) $(LDFLAGS) $^ $(LDLIBS) $(OUTPUT_FILE)
$(NVCC_CUDNN): llmc/cudnn_att.cpp
$(NVCC) -c $(NVCC_FLAGS) $(PFLAGS) $^ $(NVCC_INCLUDES) -o $@
train_gpt2cu: train_gpt2.cu $(NVCC_CUDNN)
$(NVCC) $(NVCC_FLAGS) $(PFLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE)
train_gpt2fp32cu: train_gpt2_fp32.cu
$(NVCC) $(NVCC_FLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE)
test_gpt2cu: test_gpt2.cu $(NVCC_CUDNN)
$(NVCC) $(NVCC_FLAGS) $(PFLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE)
test_gpt2fp32cu: test_gpt2_fp32.cu
$(NVCC) $(NVCC_FLAGS) $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE)
profile_gpt2cu: profile_gpt2.cu $(NVCC_CUDNN)
$(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE)
clean:
$(REMOVE_FILES) $(TARGETS)
$(REMOVE_BUILD_OBJECT_FILES)
================================================
FILE: README.md
================================================
# llm.c
LLMs in simple, pure C/CUDA with no need for 245MB of PyTorch or 107MB of cPython. Current focus is on pretraining, in particular reproducing the [GPT-2](https://github.com/openai/gpt-2) and [GPT-3](https://arxiv.org/abs/2005.14165) miniseries, along with a parallel PyTorch reference implementation in [train_gpt2.py](train_gpt2.py). You'll recognize this file as a slightly tweaked [nanoGPT](https://github.com/karpathy/nanoGPT), an earlier project of mine. Currently, llm.c is a bit faster than PyTorch Nightly (by about 7%). In addition to the bleeding edge mainline code in [train_gpt2.cu](train_gpt2.cu), we have a simple reference CPU fp32 implementation in ~1,000 lines of clean code in one file [train_gpt2.c](train_gpt2.c). I'd like this repo to only maintain C and CUDA code. Ports to other languages or repos are very welcome, but should be done in separate repos, and I am happy to link to them below in the "notable forks" section. Developer coordination happens in the [Discussions](https://github.com/karpathy/llm.c/discussions) and on Discord, either the `#llmc` channel on the [Zero to Hero](https://discord.gg/3zy8kqD9Cp) channel, or on `#llmdotc` on [GPU MODE](https://discord.gg/gpumode) Discord.
## quick start
The best introduction to the llm.c repo today is reproducing the GPT-2 (124M) model. [Discussion #481](https://github.com/karpathy/llm.c/discussions/481) steps through this in detail. We can reproduce other models from the GPT-2 and GPT-3 series in both llm.c and in the parallel implementation of PyTorch. Have a look at the [scripts README](scripts/README.md).
debugging tip: when you run the `make` command to build the binary, modify it by replacing `-O3` with `-g` so you can step through the code in your favorite IDE (e.g. vscode).
## quick start (1 GPU, fp32 only)
If you won't be training on multiple nodes, aren't interested in mixed precision, and are interested in learning CUDA, the fp32 (legacy) files might be of interest to you. These are files that were "checkpointed" early in the history of llm.c and frozen in time. They are simpler, more portable, and possibly easier to understand. Run the 1 GPU, fp32 code like this:
```bash
chmod u+x ./dev/download_starter_pack.sh
./dev/download_starter_pack.sh
make train_gpt2fp32cu
./train_gpt2fp32cu
```
The download_starter_pack.sh script is a quick & easy way to get started and it downloads a bunch of .bin files that help get you off the ground. These contain: 1) the GPT-2 124M model saved in fp32, in bfloat16, 2) a "debug state" used in unit testing (a small batch of data, and target activations and gradients), 3) the GPT-2 tokenizer, and 3) the tokenized [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset. Alternatively, instead of running the .sh script, you can re-create these artifacts manually as follows:
```bash
pip install -r requirements.txt
python dev/data/tinyshakespeare.py
python train_gpt2.py
```
## quick start (CPU)
The "I am so GPU poor that I don't even have one GPU" section. You can still enjoy seeing llm.c train! But you won't go too far. Just like the fp32 version above, the CPU version is an even earlier checkpoint in the history of llm.c, back when it was just a simple reference implementation in C. For example, instead of training from scratch, you can finetune a GPT-2 small (124M) to output Shakespeare-like text, as an example:
```bash
chmod u+x ./dev/download_starter_pack.sh
./dev/download_starter_pack.sh
make train_gpt2
OMP_NUM_THREADS=8 ./train_gpt2
```
If you'd prefer to avoid running the starter pack script, then as mentioned in the previous section you can reproduce the exact same .bin files and artifacts by running `python dev/data/tinyshakespeare.py` and then `python train_gpt2.py`.
The above lines (1) download an already tokenized [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset and download the GPT-2 (124M) weights, (3) init from them in C and train for 40 steps on tineshakespeare with AdamW (using batch size 4, context length only 64), evaluate validation loss, and sample some text. Honestly, unless you have a beefy CPU (and can crank up the number of OMP threads in the launch command), you're not going to get that far on CPU training LLMs, but it might be a good demo/reference. The output looks like this on my MacBook Pro (Apple Silicon M3 Max):
```
[GPT-2]
max_seq_len: 1024
vocab_size: 50257
num_layers: 12
num_heads: 12
channels: 768
num_parameters: 124439808
train dataset num_batches: 1192
val dataset num_batches: 128
num_activations: 73323776
val loss 5.252026
step 0: train loss 5.356189 (took 1452.121000 ms)
step 1: train loss 4.301069 (took 1288.673000 ms)
step 2: train loss 4.623322 (took 1369.394000 ms)
step 3: train loss 4.600470 (took 1290.761000 ms)
... (trunctated) ...
step 39: train loss 3.970751 (took 1323.779000 ms)
val loss 4.107781
generating:
---
Come Running Away,
Greater conquer
With the Imperial blood
the heaviest host of the gods
into this wondrous world beyond.
I will not back thee, for how sweet after birth
Netflix against repounder,
will not
flourish against the earlocks of
Allay
---
```
## datasets
The data files inside `/dev/data/(dataset).py` are responsible for downloading, tokenizing and saving the tokens to .bin files, readable easily from C. So for example when you run:
```bash
python dev/data/tinyshakespeare.py
```
We download and tokenize the [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset. The output of this looks like this:
```
writing 32,768 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_val.bin
writing 305,260 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_train.bin
```
The .bin files contain a short header (1024 bytes) and then a stream of tokens in uint16, indicating the token ids with the GPT-2 tokenizer. More datasets are available in `/dev/data`.
## test
I am also attaching a simple unit test for making sure our C code agrees with the PyTorch code. On the CPU as an example, compile and run with:
```bash
make test_gpt2
./test_gpt2
```
This now loads the `gpt2_124M_debug_state.bin` file that gets written by train_gpt2.py, runs a forward pass, compares the logits and loss with the PyTorch reference implementation, then it does 10 iterations of training with Adam and makes sure the losses match PyTorch. To test the GPU version we run:
```bash
# fp32 test (cudnn not supported)
make test_gpt2cu PRECISION=FP32 && ./test_gpt2cu
# mixed precision cudnn test
make test_gpt2cu USE_CUDNN=1 && ./test_gpt2cu
```
This tests both the fp32 path and the mixed precision path. The test should pass and print `overall okay: 1`.
## tutorial
I attached a very small tutorial here, in [doc/layernorm/layernorm.md](doc/layernorm/layernorm.md). It's a simple, step-by-step guide to implementing a single layer of the GPT-2 model, the layernorm layer. This is a good starting point to understand how the layers are implemented in C.
**flash attention**. As of May 1, 2024 we use the Flash Attention from cuDNN. Because cuDNN bloats the compile time from a few seconds to ~minute and this code path is right now very new, this is disabled by default. You can enable it by compiling like this:
```bash
make train_gpt2cu USE_CUDNN=1
```
This will try to compile with cudnn and run it. You have to have cuDNN installed on your system. The [cuDNN installation instructions](https://developer.nvidia.com/cudnn) with apt-get will grab the default set of cuDNN packages. For a minimal setup, the cuDNN dev package is sufficient, e.g. on Ubuntu 22.04 for CUDA 12.x:
```bash
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install libcudnn9-dev-cuda-12
```
On top of this you need the [cuDNN frontend](https://github.com/NVIDIA/cudnn-frontend/tree/main), but this is just header files. Simply clone the repo to your disk. The Makefile currently looks for it in either your home directory or the current directory. If you have put it elsewhere, add `CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include` to the `make` command-line.
## multi-GPU training
Make sure you install MPI and NCCL, e.g. on Linux:
```bash
sudo apt install openmpi-bin openmpi-doc libopenmpi-dev
```
For NCCL follow the instructions from the [official website](https://developer.nvidia.com/nccl/nccl-download) (e.g. network installer)
and then:
```bash
make train_gpt2cu
mpirun -np <number of GPUs> ./train_gpt2cu
```
or simply run one of our scripts under `./scripts/`.
## multi-node training
Make sure you've installed `NCCL` following instructions from [multi-GPU](#multi-gpu-training) section.
There are 3 ways we currently support that allow you to run multi-node training:
1) Use OpenMPI to exchange nccl id and initialize NCCL. See e.g. `./scripts/multi_node/run_gpt2_124M_mpi.sh` script for details.
2) Use shared file system to init NCCL. See `./scripts/multi_node/run_gpt2_124M_fs.sbatch` script for details.
3) Use TCP sockets to init NCCL. See `./scripts/multi_node/run_gpt2_124M_tcp.sbatch` script for details.
Note:
* If you're running in a slurm environment and your slurm doesn't support PMIx (which we assume will be a common situation given that `slurm-wlm` dropped PMIx support) you will have to use FS (2) or TCP (3) approach. To test whether your slurm supports PMIx run: `srun --mpi=list` and see whether you get `pmix` in the output.
* If you don't have slurm set up, you can kick off a multi-node run using `mpirun` - MPI (1).
None of these 3 methods is superior, we just offer you options so that you can run in your specific environment.
## experiments / sweeps
Just as an example process to sweep learning rates on a machine with 4 GPUs on TinyStories. Run a shell script `sweep.sh` (after you of course `chmod u+x sweep.sh`):
```bash
#!/bin/bash
learning_rates=(3e-5 1e-4 3e-4 1e-3)
for i in {0..3}; do
export CUDA_VISIBLE_DEVICES=$i
screen -dmS "tr$i" bash -c "./train_gpt2cu -i data/TinyStories -v 250 -s 250 -g 144 -l ${learning_rates[$i]} -o stories$i.log"
done
# you can bring these down with
# screen -ls | grep -E "tr[0-3]" | cut -d. -f1 | xargs -I {} screen -X -S {} quit
```
This example opens up 4 screen sessions and runs the four commands with different LRs. This writes the log files `stories$i.log` with all the losses, which you can plot as you wish in Python. A quick example of how to parse and plot these logfiles is in [dev/vislog.ipynb](dev/vislog.ipynb).
## repo
A few more words on what I want this repo to be:
First, I want `llm.c` to be a place for education. E.g. our `dev/cuda` folder is a place for a library of kernels for all the layers that are manually hand-written and very well documented, starting from very simple kernels all the way to more complex / faster kernels. If you have a new kernel with various different tradeoffs, please feel free to contribute it here.
That said, I also want `llm.c` to be very fast too, even practically useful to train networks. E.g. to start, we should be able to reproduce the big GPT-2 (1.6B) training run. This requires that we incorporate whatever fastest kernels there are, including the use of libraries such as cuBLAS, cuBLASLt, CUTLASS, cuDNN, etc. I also think doing so serves an educational purpose to establish an expert upper bound, and a unit of measurement, e.g. you could say that your manually written kernels are 80% of cuBLAS speed, etc. Then you can choose to do a super fast run, or you can choose to "drag and drop" whatever manual kernels you wish to use, and run with those.
However, as a constraint, I want to keep the mainline `llm.c` in the root folder simple and readable. If there is a PR that e.g. improves performance by 2% but it "costs" 500 lines of complex C code, and maybe an exotic 3rd party dependency, I may reject the PR because the complexity is not worth it. As a concrete example - making cuBLAS for matmuls the default in the root training loop is a no-brainer: it makes the mainline code much faster, it is a single line of interpretable code, and it is a very common dependency. On the side of this, we can have manual implementations that can compete with cuBLAS in `dev/cuda`.
Lastly, I will be a lot more sensitive to complexity in the root folder of the project, which contains the main / default files of the project. In comparison, the `dev/` folder is a bit more of a scratch space for us to develop a library of kernels or classes and share useful or related or educational code, and some of this code could be ok to be (locally) complex.
## notable forks
- AMD support
- [llm.c](https://github.com/anthonix/llm.c) by @[anthonix](https://github.com/anthonix): support for AMD devices, such as the 7900 XTX
- C#
- [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): a C# port of this project
- [Llm.cs](https://github.com/nietras/Llm.cs) by @[nietras](https://github.com/nietras): a C# port of this project with focus on easy to get started on any platform. Clone and run ✅
- CUDA C++
- [llm.cpp](https://github.com/gevtushenko/llm.c) by @[gevtushenko](https://github.com/gevtushenko): a port of this project using the [CUDA C++ Core Libraries](https://github.com/NVIDIA/cccl)
- A presentation this fork was covered in [this lecture](https://www.youtube.com/watch?v=WiB_3Csfj_Q) in the [GPU MODE Discord Server](https://discord.gg/cudamode)
- C++/CUDA
- [llm.cpp](https://github.com/zhangpiu/llm.cpp/tree/master/llmcpp) by @[zhangpiu](https://github.com/zhangpiu): a port of this project using the [Eigen](https://gitlab.com/libeigen/eigen), supporting CPU/CUDA.
- WebGPU C++
- [gpu.cpp](https://github.com/AnswerDotAI/gpu.cpp) by @[austinvhuang](https://github.com/austinvhuang): a library for portable GPU compute in C++ using native WebGPU. Aims to be a general-purpose library, but also porting llm.c kernels to WGSL.
- C++
- [llm.cpp](https://github.com/GaoYusong/llm.cpp) by @[GaoYusong](https://github.com/GaoYusong): a port of this project featuring a C++ single-header [tinytorch.hpp](https://github.com/GaoYusong/llm.cpp/blob/main/tinytorch.hpp) library
- Go
- [llm.go](https://github.com/joshcarp/llm.go) by @[joshcarp](https://github.com/joshcarp): a Go port of this project
- Java
- [llm.java](https://github.com/harryjackson/llm.java) by @[harryjackson](https://github.com/harryjackson): a Java port of this project
- Metal
- [llm.metal](https://github.com/regrettable-username/llm.metal) by @[regrettable-username](https://github.com/regrettable-username): LLM training in simple, raw C/Metal Shading Language
- Mojo
- [llm.🔥](https://github.com/dorjeduck/llm.mojo) by @[dorjeduck](https://github.com/dorjeduck): a Mojo port of this project
- OpenCL
- [llm.c](https://github.com/krrishnarraj/llm.c) by @[krrishnarraj](https://github.com/krrishnarraj): an OpenCL port of this project
- Rust
- [llm.rs](https://github.com/yijunyu/llm.rs) by @[Yijun Yu](https://github.com/yijunyu): a Rust rewrite with the aim to have same performance
- [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project
- Swift
- [llm.swift](https://github.com/otabuzzman/llm.swift) by @[otabuzzman](https://github.com/otabuzzman): a Swift port of this project
- Zig
- [llm.zig](https://github.com/Saimirbaci/llm.zig) by @[saimirbaci](https://github.com/Saimirbaci): a Zig port of this project
- Habana Gaudi2
- [llm.tpc](https://github.com/abhilash1910/llm.tpc) by @[abhilash1910](https://github.com/abhilash1910): a Habana Gaudi2 port of this project
- Nim
- [llm.nim](https://github.com/Vindaar/llm.nim) by @[Vindaar](https://github.com/Vindaar): a Nim port of this project
## discussions
Ways of organizing development:
- Experiencing a concrete issue with the repo? Use [Issues](https://github.com/karpathy/llm.c/issues).
- Have some code to contribute? Open a [PR](https://github.com/karpathy/llm.c/pulls)
- Chat about the repo, ask questions, etc.? Look at [Discussions](https://github.com/karpathy/llm.c/discussions).
- Something faster? I created a new `#llmc` channel on my [Zero to Hero Discord channel](https://discord.gg/3zy8kqD9Cp).
## license
MIT
================================================
FILE: dev/cpu/matmul_forward.c
================================================
/*
CPU Kernels for matmul forward pass.
*/
// Compile Examples:
//
// MSVC: cl.exe /O2 /fp:fast /Qvec-report:2 /I. /I ..\..\dev matmul_forward.c
// cl.exe /O2 /fp:fast /Qvec-report:2 /arch:AVX /I. /I ..\..\dev matmul_forward.c
// cl.exe /O2 /fp:fast /Qvec-report:2 /arch:AVX2 /I. /I ..\..\dev matmul_forward.c
//
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <unistd.h>
// ----------------------------------------------------------------------------
// CPU code reference
void matmul_forward_cpu(float* out,
const float* inp, const float* weight, const float* bias,
int B, int T, int C, int OC) {
// OC is short for "output channels"
// inp is (B,T,C), weight is (OC, C), bias is (OC)
// out will be (B,T,OC)
for (int b = 0; b < B; b++) {
for (int t = 0; t < T; t++) {
float* out_bt = out + b * T * OC + t * OC;
const float* inp_bt = inp + b * T * C + t * C;
for (int o = 0; o < OC; o++) {
float val = (bias != NULL) ? bias[o] : 0.0f;
const float* wrow = weight + o*C;
for (int i = 0; i < C; i++) {
val += inp_bt[i] * wrow[i];
}
out_bt[o] = val;
}
}
}
}
void matmul_forward_ngc92(float* out,
const float* inp, const float* weight, const float* bias,
int B, int T, int C, int OC) {
// most of the running time is spent here and in matmul_backward
// OC is short for "output channels"
// inp is (B,T,C), weight is (OC, C), bias is (OC)
// out will be (B,T,OC)
// make sure the tiled loop will be correct, otherwise, fallback to slow version
#define LOOP_UNROLL 8
if (B * T % LOOP_UNROLL != 0) {
printf("MUST BE A MULTIPLE OF 8"); // FIXME
return;
}
// collapse the B and T loops into one and turn it into a strided loop.
// then we can tile the inner loop, and reuse the loaded weight LOOP_UNROLL many times
// for significant speed-ups.
for (int obt = 0; obt < B * T; obt += LOOP_UNROLL) {
for (int o = 0; o < OC; o++) {
// keep LOOP_UNROLL many results in register, initialized by the bias term.
float result[LOOP_UNROLL];
for (int ibt = 0; ibt < LOOP_UNROLL; ++ibt) {
result[ibt] = (bias != NULL) ? bias[o] : 0.0f;
}
// inner loops. Because we do LOOP_UNROLL steps of inner bt, we can cache
// the value of weight[i + o * C] and reuse it.
// we compile with -Ofast, so the compiler will turn the inner loop into a bunch of FMAs
for (int i = 0; i < C; i++) {
float w = weight[i + o * C];
for (int ibt = 0; ibt < LOOP_UNROLL; ++ibt) {
int bt = obt + ibt;
result[ibt] += inp[bt * C + i] * w;
}
}
// write back results to main memory
for (int ibt = 0; ibt < LOOP_UNROLL; ++ibt) {
int bt = obt + ibt;
out[bt * OC + o] = result[ibt];
}
}
}
}
#define NUM_KERNELS 2
void matmul_forward(int kernel_num,
float* out,
const float* inp, const float* weight, const float* bias,
int B, int T, int C, int OC) {
switch (kernel_num) {
case 0:
matmul_forward_cpu(out, inp, weight, bias, B, T, C, OC);
break;
case 1:
matmul_forward_ngc92(out, inp, weight, bias, B, T, C, OC);
break;
default:
printf("Invalid kernel number\n");
exit(1);
}
}
void validate_results_cpu(const float* device_result, const float* cpu_reference, const char* name, int num_elements, float tolerance);
float* make_random_float(size_t N);
int main(int argc, char **argv) {
srand(0);
int B = 8;
int T = 1024;
int C = 768;
int OC = 768 * 4; // expansion of 4, e.g. in the MLP
int RUNS = 4; // number of times to run a kernel for benchmarks
srand(137);
float* out = make_random_float(B * T * OC);
float* inp = make_random_float(B * T * C);
float* weight = make_random_float(OC * C);
float* bias = make_random_float(OC);
float* grad_out = make_random_float(B * T * OC);
float* grad_inp = make_random_float(B * T * C);
float* grad_weight = make_random_float(OC * C);
float* grad_bias = make_random_float(OC);
printf("> Calculating reference\n");
matmul_forward_cpu(out, inp, weight, bias, B, T, C, OC);
for (int kernel_num = 0; kernel_num < NUM_KERNELS; kernel_num++) {
printf("> Verifying kernel #%d\n", kernel_num);
srand(137);
float* kernel_out = make_random_float(B * T * OC);
float* kernel_inp = make_random_float(B * T * C);
float* kernel_weight = make_random_float(OC * C);
float* kernel_bias = make_random_float(OC);
matmul_forward(kernel_num, kernel_out, kernel_inp, kernel_weight, kernel_bias, B, T, C, OC);
validate_results_cpu(kernel_out, out, "out", B * T * OC, 1e-5);
free(kernel_out);
free(kernel_inp);
free(kernel_weight);
free(kernel_bias);
}
printf("All kernels passed! Starting benchmarks.\n\n");
for (int kernel_num = 0; kernel_num < NUM_KERNELS; kernel_num++) {
printf("> Running kernel #%d\n", kernel_num);
struct timespec start, end;
clock_gettime(CLOCK_MONOTONIC, &start);
for (int i = 0; i < RUNS; i++) {
matmul_forward(kernel_num, out, inp, weight, bias, B, T, C, OC);
}
clock_gettime(CLOCK_MONOTONIC, &end);
double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
printf("> Kernel #%d, (took %f ms)\n", kernel_num, time_elapsed_s * 1000);
}
// free memory
free(out);
free(inp);
free(weight);
free(bias);
free(grad_out);
free(grad_inp);
free(grad_weight);
free(grad_bias);
return 0;
}
float* make_random_float(size_t N) {
float* arr = (float*)malloc(N * sizeof(float));
for (size_t i = 0; i < N; i++) {
arr[i] = ((float)rand() / RAND_MAX) * 2.0 - 1.0; // range -1..1
}
return arr;
}
void validate_results_cpu(const float* kernel_result, const float* cpu_reference, const char* name, int num_elements, float tolerance) {
int nfaults = 0;
for (int i = 0; i < num_elements; i++) {
// print the first few comparisons
if (i < 5) {
printf("%f %f\n", cpu_reference[i], kernel_result[i]);
}
float t_eff = tolerance + fabs(cpu_reference[i]);
// ensure correctness for all elements.
if (fabs(cpu_reference[i] - kernel_result[i]) > t_eff) {
printf("Mismatch of %s at %d: CPU_ref: %f vs CPU_new: %f\n", name, i, cpu_reference[i], kernel_result[i]);
nfaults++;
if (nfaults >= 10) {
exit(EXIT_FAILURE);
}
}
}
if (nfaults > 0) {
exit(EXIT_FAILURE);
}
printf("OK\n");
}
================================================
FILE: dev/cuda/Makefile
================================================
# Makefile for building dev/cuda kernels
# Collects all the make commands in one file but each file also
# has the compile and run commands in the header comments section.
# Find nvcc (NVIDIA CUDA compiler)
NVCC := $(shell which nvcc 2>/dev/null)
ifeq ($(NVCC),)
$(error nvcc not found.)
endif
ifneq ($(CI),true) # if not in CI, then use the GPU query
ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY=
GPU_COMPUTE_CAPABILITY = $(shell __nvcc_device_query) # assume if NVCC is present, then this likely is too
GPU_COMPUTE_CAPABILITY := $(strip $(GPU_COMPUTE_CAPABILITY))
endif
endif
# Compiler flags
ifeq ($(GPU_COMPUTE_CAPABILITY),) # set to defaults if: make GPU_COMPUTE_CAPABILITY=
CFLAGS = -O3 --use_fast_math
else
CFLAGS = -O3 --use_fast_math --generate-code arch=compute_$(GPU_COMPUTE_CAPABILITY),code=[compute_$(GPU_COMPUTE_CAPABILITY),sm_$(GPU_COMPUTE_CAPABILITY)]
endif
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
%: %.cu
$(NVCC) $(CFLAGS) $(NVCCFLAGS) $< -o $@
# Build all targets
TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward fused_residual_forward global_norm permute
all: $(TARGETS)
all_ptx: $(TARGETS:%=%.ptx)
all_sass: $(TARGETS:%=%.sass)
# Individual targets: forward pass
attention_forward: attention_forward.cu
classifier_fused: classifier_fused.cu
crossentropy_forward: crossentropy_forward.cu
encoder_forward: encoder_forward.cu
gelu_forward: gelu_forward.cu
layernorm_forward: layernorm_forward.cu
fused_residual_forward: fused_residual_forward.cu
residual_forward: residual_forward.cu
softmax_forward: softmax_forward.cu
trimat_forward: trimat_forward.cu
# matmul fwd/bwd also uses OpenMP (optionally) and cuBLASLt libs
matmul_forward: matmul_forward.cu
$(NVCC) $(CFLAGS) $(NVCCFLAGS) -Xcompiler -fopenmp matmul_forward.cu -o matmul_forward
# Individual targets: backward pass
attention_backward: attention_backward.cu
crossentropy_softmax_backward: crossentropy_softmax_backward.cu
encoder_backward: encoder_backward.cu
gelu_backward: gelu_backward.cu
layernorm_backward: layernorm_backward.cu
matmul_backward_bias: matmul_backward_bias.cu
matmul_backward: matmul_backward.cu
$(NVCC) $(CFLAGS) $(NVCCFLAGS) -Xcompiler -fopenmp matmul_backward.cu -o matmul_backward
# Update kernels
adamw: adamw.cu
global_norm: global_norm.cu
permute: permute.cu
# NCCL communication kernels
nccl_all_reduce: nccl_all_reduce.cu
$(NVCC) -lmpi -lnccl $(NVCCFLAGS) $(MPI_PATHS) nccl_all_reduce.cu -o nccl_all_reduce
# Generate PTX using cuobjdump
%.ptx: %
cuobjdump --dump-ptx $< > $@
# Generate SASS using cuobjdump
%.sass: %
cuobjdump --dump-sass $< > $@
# Run all targets
run_all: all
@for target in $(TARGETS); do \
echo "\n========================================"; \
echo "Running $$target ..."; \
echo "========================================\n"; \
./$$target; \
done
# Clean up
clean:
rm -f $(TARGETS) *.ptx *.sass
================================================
FILE: dev/cuda/README.md
================================================
# dev/cuda
This directory is scratch space for developing various versions of the needed CUDA kernels. Each file develops a kernel, and usually multiple versions of that kernel that could have different running times and of different code or time complexity.
See the top of each file for how to compile and run the kernel. Alternatively, the commands are also all grouped in the `Makefile` in this directory for convenience.
For example, we can look at the top of `layernorm_forward.cu` to build the forward pass kernels for the LayerNorm:
```bash
nvcc -O3 --use_fast_math -lcublas -lcublasLt layernorm_forward.cu -o layernorm_forward
```
or simply
```bash
make layernorm_forward
```
The comments at the top then document the different versions of this kernel available, usually these are in increasing complexity and decreasing running times. For example, inspecting the comments in the file on top, the most naive kernel we can then run as:
```bash
./layernorm_forward 1
```
You'll see that this first forwards the reference code on the CPU, then it runs kernel 1 on the GPU, compares the results to check for correctness, and then runs a number of configurations of this kernel (most often and most notably the block size), to time the kernel in these launch configurations. We can then run one of the faster kernels (kernel 4) instead:
```bash
./layernorm_forward 4
```
You'll see that this matches all the CPU results but runs much much faster. The typical process from here on is we copy paste the kernel that ran fastest, adjust it manually (e.g. to hardcode the best block size) and drop it into the training code file, e.g. `train_gpt2.cu`.
To add a new version of a kernel, add the kernel to the corresponding file and adjust the docs. To add a new kernel, add the new file and adjust the Makefile. Run `make clean` to clean up binaries from your directory.
If you do not have a GPU or is having trouble with CUDA dependencies, you can run the benchmarks on the [Modal platform](http://modal.com). For example, to run the benchmark for the attention forward pass on an A100 GPU with 80GB of memory, you can run the following command:
```bash
GPU_MEM=80 modal run benchmark_on_modal.py --compile-command "nvcc -O3 --use_fast_math attention_forward.cu -o attention_forward -lcublas" --run-command "./attention_forward 1"
```
================================================
FILE: dev/cuda/adamw.cu
================================================
/*
Kernels for the AdamW optimizer.
References:
* https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html
* https://github.com/nvidia/apex/blob/master/csrc/multi_tensor_adam.cu
Compile example:
nvcc -lcublas -lcublasLt adamw.cu -o adamw
nvcc -O3 --use_fast_math -lcublas -lcublasLt adamw.cu -o adamw
./adamw
TODO(general):
amsgrad=True
TODO(perf):
dtype
thread coarsening/ILP
*/
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <cuda_runtime.h>
#include "common.h"
// ----------------------------------------------------------------------------
// CPU code reference
void adamw_cpu(float* params_memory, const float* grads_memory, float* m_memory, float* v_memory, int t, long num_parameters, float learning_rate=1e-3, float beta1=0.9, float beta2=0.999, float eps=1e-8, float weight_decay=0.0) {
// adapted from: train_gpt2.c
for (int i = 0; i < num_parameters; i++) {
float param = params_memory[i];
float grad = grads_memory[i];
// update the first moment (momentum)
float m = beta1 * m_memory[i] + (1.0f - beta1) * grad;
// update the second moment (RMSprop)
float v = beta2 * v_memory[i] + (1.0f - beta2) * grad * grad;
// bias-correct both moments
float m_hat = m / (1.0f - powf(beta1, t));
float v_hat = v / (1.0f - powf(beta2, t));
// update
m_memory[i] = m;
v_memory[i] = v;
params_memory[i] -= learning_rate * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * param);
}
}
// ----------------------------------------------------------------------------
// GPU kernels
// utility functions
// 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));
}
// naive fused kernel
__global__ void adamw_kernel1(float* params_memory, const 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
// update the first moment (momentum)
m_memory[i] = beta1 * m_memory[i] + (1.0f - beta1) * grads_memory[i];
// update the second moment (RMSprop)
v_memory[i] = beta2 * v_memory[i] + (1.0f - beta2) * grads_memory[i] * grads_memory[i];
float m_hat = m_memory[i] / beta1_correction;
float v_hat = v_memory[i] / beta2_correction;
params_memory[i] -= learning_rate * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * params_memory[i]);
}
// Slightly more optimized AdamW kernel by:
// * loading data that is accessed more than once into registers,
// * using optimized linear interpolation for the moment updates.
__global__ void adamw_kernel2(float* params_memory, const 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]);
}
// ----------------------------------------------------------------------------
// kernel launcher
// version 1: naive dispatch to naive kernel
void adamw_dispatch1(float* params_memory, const 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) {
unsigned int block_size = 512;
unsigned int num_blocks = ceil_div(num_parameters, (long) block_size);
adamw_kernel1<<<num_blocks, block_size>>>(params_memory, grads_memory, m_memory, v_memory, num_parameters,
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);
cudaCheck(cudaGetLastError());
}
// version 2: naive dispatch to slightly optimized kernel
void adamw_dispatch2(float* params_memory, const 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) {
unsigned int block_size = 512;
unsigned int num_blocks = ceil_div(num_parameters, (long) block_size);
adamw_kernel2<<<num_blocks, block_size>>>(params_memory, grads_memory, m_memory, v_memory, num_parameters,
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);
cudaCheck(cudaGetLastError());
}
void adamw(int kernel_num,
float* params_memory, const float* grads_memory, float* m_memory, float* v_memory, int t, long num_parameters,
float learning_rate=1e-3, float beta1=0.9, float beta2=0.999, float eps=1e-8, float weight_decay=0.0) {
// calculate the m_hat and v_hat correction terms once as they are the same for every param/thread
float beta1_correction = 1.0f - powf(beta1, t);
float beta2_correction = 1.0f - powf(beta2, t);
switch (kernel_num) {
case 1:
adamw_dispatch1(params_memory, grads_memory, m_memory, v_memory, num_parameters,
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);
break;
case 2:
adamw_dispatch2(params_memory, grads_memory, m_memory, v_memory, num_parameters,
learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay);
break;
default:
printf("Invalid kernel number\n");
exit(1);
}
}
// ----------------------------------------------------------------------------
int main(int argc, char **argv) {
setup_main();
const long num_parameters = 1048576;
const int t = 10;
const float learning_rate = 1e-3f;
const float beta1 = 0.9f;
const float beta2 = 0.999f;
const float eps = 1e-8f;
const float weight_decay = 0.0f;
// create random data on host (to be used for the CPU reference implementation)
float* params_memory = make_random_float(num_parameters);
float* grads_memory = make_random_float(num_parameters);
float* m_memory = make_random_float(num_parameters);
float* v_memory = make_random_float_01(num_parameters);
// move to GPU
float* d_params_memory;
float* d_grads_memory;
float* d_m_memory;
float* d_v_memory;
cudaCheck(cudaMalloc(&d_params_memory, num_parameters * sizeof(float)));
cudaCheck(cudaMalloc(&d_grads_memory, num_parameters * sizeof(float)));
cudaCheck(cudaMalloc(&d_m_memory, num_parameters * sizeof(float)));
cudaCheck(cudaMalloc(&d_v_memory, num_parameters * sizeof(float)));
cudaCheck(cudaMemcpy(d_params_memory, params_memory, num_parameters * sizeof(float), cudaMemcpyHostToDevice));
cudaCheck(cudaMemcpy(d_grads_memory, grads_memory, num_parameters * sizeof(float), cudaMemcpyHostToDevice));
cudaCheck(cudaMemcpy(d_m_memory, m_memory, num_parameters * sizeof(float), cudaMemcpyHostToDevice));
cudaCheck(cudaMemcpy(d_v_memory, v_memory, num_parameters * sizeof(float), cudaMemcpyHostToDevice));
// read kernel_num from command line
int kernel_num = 1;
if (argc > 1) {
kernel_num = atoi(argv[1]);
}
printf("Using kernel %d\n", kernel_num);
// calculate the CPU reference (using default hyperparams)
clock_t start = clock();
adamw_cpu(params_memory, grads_memory, m_memory, v_memory, t, num_parameters);
clock_t end = clock();
// TODO: measure runtime with multiple runs
double elapsed_time_cpu = (double)(end - start) / CLOCKS_PER_SEC;
// calculate the GPU version (using default hyperparams)
adamw(kernel_num, d_params_memory, d_grads_memory, d_m_memory, d_v_memory, t, num_parameters);
// compare
printf("Checking correctness...\n");
printf("parameters:\n");
validate_result(d_params_memory, params_memory, "params_memory", num_parameters);
printf("first moment:\n");
validate_result(d_m_memory, m_memory, "m_memory", num_parameters);
printf("second moment:\n");
validate_result(d_v_memory, v_memory, "v_memory", num_parameters);
printf("All results match.\n\n");
// now benchmark the kernel
int repeat_times = 1000;
float elapsed_time = benchmark_kernel(repeat_times, adamw, kernel_num,
d_params_memory, d_grads_memory, d_m_memory, d_v_memory, t, num_parameters,
learning_rate, beta1, beta2, eps, weight_decay);
printf("time gpu %.4f ms\n", elapsed_time);
printf("time cpu %.4f ms\n", elapsed_time_cpu);
// cleanup
free(params_memory);
free(grads_memory);
free(m_memory);
free(v_memory);
cudaCheck(cudaFree(d_params_memory));
cudaCheck(cudaFree(d_grads_memory));
cudaCheck(cudaFree(d_m_memory));
cudaCheck(cudaFree(d_v_memory));
return 0;
}
================================================
FILE: dev/cuda/attention_backward.cu
================================================
/*
Kernels for attention backward pass.
Compile example:
nvcc -O3 --use_fast_math -lcublas -lcublasLt attention_backward.cu -o attention_backward
version 1 is a naive first version
OMP_NUM_THREADS=32 ./attention_backward 1
version 2 much ensures better load-balancing by having independent threads for each batch and attention head
OMP_NUM_THREADS=32 ./attention_backward 2
version 3 uses a full warp to calculate each result (instead of a thread), which enables coalesced memory access
OMP_NUM_THREADS=32 ./attention_backward 3
version 4 improves data reuse in registers by doing 8 values of t3 in one warp.
OMP_NUM_THREADS=32 ./attention_backward 4
version 5 reduces the amount of non-fp32 instructions needed by avoiding ifs
OMP_NUM_THREADS=32 ./attention_backward 5
*/
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <float.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cooperative_groups/scan.h>
#include "common.h"
// ----------------------------------------------------------------------------
// CPU code reference
/*
NOTE:
This version of attention_forward is modified to be consistent with the
attention_forward GPU kernel in the following way small but important way:
- preatt is only QUERY @ KEY, without the scale
- the scale instead moved and fused into the softmax
- the full preatt matrix is materialized, even the parts that get masked out
- this doesn't actually change anything due to masking, but it lets us
easily compare to the GPU version, which also does the full, dense sgemm
In this way we'll be able to make sure that preatt and att agree CPU vs GPU
*/
void attention_forward_cpu(float* out, float* preatt, float* att,
float* inp,
int B, int T, int C, int NH) {
// input is (B, T, 3C) holding the query, key, value (Q, K, V) vectors
// preatt, att are (B, NH, T, T). NH = number of heads, T = sequence length
// that holds the pre-attention and post-attention scores (used in backward)
// output is (B, T, C)
// attention is the only layer that mixes information across time
// every other operation is applied at every (b,t) position independently
// (and of course, no layer mixes information across batch)
int C3 = C*3;
int hs = C / NH; // head size
float scale = 1.0 / sqrtf(hs);
#pragma omp parallel for collapse(3)
for (int b = 0; b < B; b++) {
for (int t = 0; t < T; t++) {
for (int h = 0; h < NH; h++) {
float* query_t = inp + b * T * C3 + t * C3 + h * hs;
float* preatt_bth = preatt + b*NH*T*T + h*T*T + t*T;
float* att_bth = att + b*NH*T*T + h*T*T + t*T;
// pass 1: calculate query dot key and maxval
float maxval = -FLT_MAX;
for (int t2 = 0; t2 < T; t2++) { // used to be t2 <= t
float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key
// (query_t) dot (key_t2)
float val = 0.0f;
for (int i = 0; i < hs; i++) {
val += query_t[i] * key_t2[i];
}
if (val > maxval) {
maxval = val;
}
preatt_bth[t2] = val;
}
// pass 2: calculate the exp and keep track of sum
// maxval is being calculated and subtracted only for numerical stability
float expsum = 0.0f;
for (int t2 = 0; t2 <= t; t2++) {
float expv = expf(scale * (preatt_bth[t2] - maxval));
expsum += expv;
att_bth[t2] = expv;
}
float expsum_inv = expsum == 0.0f ? 0.0f : 1.0f / expsum;
// pass 3: normalize to get the softmax
for (int t2 = 0; t2 < T; t2++) {
if (t2 <= t) {
att_bth[t2] *= expsum_inv;
} else {
// causal attention mask. not strictly necessary to set to zero here
// only doing this explicitly for debugging and checking to PyTorch
att_bth[t2] = 0.0f;
}
}
// pass 4: accumulate weighted values into the output of attention
float* out_bth = out + b * T * C + t * C + h * hs;
for (int i = 0; i < hs; i++) { out_bth[i] = 0.0f; }
for (int t2 = 0; t2 <= t; t2++) {
float* value_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C*2; // +C*2 because it's value
float att_btht2 = att_bth[t2];
for (int i = 0; i < hs; i++) {
out_bth[i] += att_btht2 * value_t2[i];
}
}
}
}
}
}
// NOTE: Also contains the re-shuffling of the exact position of "scale"
// and when it is applied (after preatt, not "during" preatt)
// also, full matrices are materialized, even the parts that get masked out
void attention_backward_cpu(float* dinp, float* dpreatt, float* datt,
float* dout, float* inp, float* att,
int B, int T, int C, int NH) {
// inp/dinp are (B, T, 3C) Q,K,V
// att/datt/dpreatt are (B, NH, T, T)
// dout is (B, T, C)
int C3 = C*3;
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++) {
float* att_bth = att + b*NH*T*T + h*T*T + t*T;
float* datt_bth = datt + b*NH*T*T + h*T*T + t*T;
float* dpreatt_bth = dpreatt + b*NH*T*T + h*T*T + t*T;
float* dquery_t = dinp + b * T * C3 + t * C3 + h * hs;
float* query_t = inp + b * T * C3 + t * C3 + h * hs;
// backward pass 4, through the value accumulation
float* dout_bth = dout + b * T * C + t * C + h * hs;
for (int t2 = 0; t2 < T; t2++) { // ADJUSTED! this was t2 <= t (see note on function)
float* value_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C*2; // +C*2 because it's value
float* dvalue_t2 = dinp + b * T * C3 + t2 * C3 + h * hs + C*2;
for (int i = 0; i < hs; i++) {
// in the forward pass this was:
// out_bth[i] += att_bth[t2] * value_t2[i];
// so now we have:
datt_bth[t2] += value_t2[i] * dout_bth[i];
dvalue_t2[i] += att_bth[t2] * dout_bth[i];
}
}
// backward pass 2 & 3, the softmax
// note that softmax (like e.g. tanh) doesn't need the input (preatt) to backward
for (int t2 = 0; t2 <= t; t2++) {
for (int t3 = 0; t3 <= t; t3++) {
float indicator = t2 == t3 ? 1.0f : 0.0f;
float local_derivative = att_bth[t2] * (indicator - att_bth[t3]);
dpreatt_bth[t3] += scale * local_derivative * datt_bth[t2];
}
}
// backward pass 1, the query @ key matmul
for (int t2 = 0; t2 <= t; t2++) {
float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key
float* dkey_t2 = dinp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key
for (int i = 0; i < hs; i++) {
// in the forward pass this was:
// preatt_bth[t2] += query_t[i] * key_t2[i]
// so now we have:
dquery_t[i] += key_t2[i] * dpreatt_bth[t2];
dkey_t2[i] += query_t[i] * dpreatt_bth[t2];
}
}
}
}
}
}
// ----------------------------------------------------------------------------
// GPU kernels
// the forward pass that is the sequence [permute, sgemm, softmax, sgemm, unpermute]
__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] = inp[inp_idx];
k[idx] = inp[inp_idx + NH * d];
v[idx] = 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(const 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] = 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<float*>(&vec)[index];
}
__device__ float vec_at(const float4& vec, int index) {
return reinterpret_cast<const float*>(&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);
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();
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<const float4*>(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<float>{});
sumval *= expf(inv_temperature * (maxval - global_maxval));
float sum = cg::reduce(warp, sumval, cg::plus<float>{});
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);
}
}
// naive kernel to backward through an autoregressive softmax, just to get correctness
__global__ void softmax_autoregressive_backward_kernel1(float* dpreatt, const float* datt, const float* att,
int B, int T, int C, int NH) {
// dpreatt, datt, att are all (B, NH, T, T)
int t3 = blockIdx.x * blockDim.x + threadIdx.x;
if (t3 < T) {
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
for (int b = 0; b < B; b++) {
for (int h = 0; h < NH; h++) {
for (int t = t3; t < T; t++) {
const float* att_bth = att + b*NH*T*T + h*T*T + t*T;
const float* datt_bth = datt + b*NH*T*T + h*T*T + t*T;
float* dpreatt_bth = dpreatt + b*NH*T*T + h*T*T + t*T;
float accum = 0.0f;
for (int t2 = 0; t2 <= t; t2++) {
float indicator = t2 == t3 ? 1.0f : 0.0f;
float local_derivative = att_bth[t2] * (indicator - att_bth[t3]);
accum += scale * local_derivative * datt_bth[t2];
}
dpreatt_bth[t3] = accum;
}
}
}
}
}
// parallelize across t,b,h
__global__ void softmax_autoregressive_backward_kernel2(float* dpreatt, const float* datt, const float* att,
int B, int T, int C, int NH) {
int t3 = blockIdx.x * blockDim.x + threadIdx.x;
int idx = blockIdx.y * T * T;
if (t3 >= T) { return; }
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
for (int t = t3; t < T; t++) {
float result = 0.0;
const float* att_bth = att + idx + t*T;
const float* datt_bth = datt + idx + t*T;
float* dpreatt_bth = dpreatt + idx + t*T;
for (int t2 = 0; t2 <= t; t2++) {
float indicator = t2 == t3 ? 1.0f : 0.0f;
float local_derivative = att_bth[t2] * (indicator - att_bth[t3]);
result += scale * local_derivative * datt_bth[t2];
}
dpreatt_bth[t3] = result;
}
}
// parallelize across t,b,h
__global__ void softmax_autoregressive_backward_kernel3(float* dpreatt, const float* datt, const float* att,
int B, int T, int C, int NH) {
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
int t3 = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
int idx = blockIdx.y * T * T;
if (t3 >= T) { return; }
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
for (int t = t3; t < T; t++) {
float result = 0.0;
const float* att_bth = att + idx + t*T;
const float* datt_bth = datt + idx + t*T;
float* dpreatt_bth = dpreatt + idx + t*T;
const float att_at_t3 = att_bth[t3];
for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {
float indicator = t2 == t3 ? 1.0f : 0.0f;
float local_derivative = att_bth[t2] * (indicator - att_at_t3);
result += local_derivative * datt_bth[t2];
}
result = cg::reduce(warp, result, cg::plus<float>());
if(warp.thread_rank() == 0) {
dpreatt_bth[t3] = scale * result;
}
}
}
__global__ void softmax_autoregressive_backward_kernel4(float* __restrict__ dpreatt, const float* __restrict__ datt,
const float* __restrict__ att,
int B, int T, int C, int NH) {
constexpr int UNROLL = 8;
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
int t3 = UNROLL * (blockIdx.x * warp.meta_group_size() + warp.meta_group_rank());
int idx = blockIdx.y * T * T;
if (t3 >= T) { return; }
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
// the innermost loop combines different values of t2 with different values of t.
// by handling [t3, t3 + UNROLL) in one thread, we get much better memory reuse:
// any t3/t-dependent value can be loaded once before the t2 loop.
// within the t2 loop, we can combine each loaded value with each of the UNROLL
// pre-loaded values, thus cutting memory ready by a factor of ~UNROLL.
// one iteration of this loop has to handle the cases
// this may lead to some invalid indices; therefore, we have several
// early-outs in the iteration over k below.
for (int t = t3; t < T; t++) {
float result[UNROLL] = {};
const float* att_bth = att + idx + t * T;
const float* datt_bth = datt + idx + t * T;
float* dpreatt_bth = dpreatt + idx + t * T;
float att_at_t3[UNROLL];
for(int k = 0; k < UNROLL; ++k) {
if (t < t3 + k) continue;
att_at_t3[k] = att_bth[t3 + k];
}
for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {
float att_t2 = att_bth[t2];
float datt_t2 = datt_bth[t2];
for(int k = 0; k < UNROLL; ++k) {
if (t < t3 + k) continue;
float indicator = t2 == (t3 + k) ? 1.0f : 0.0f;
float local_derivative = att_t2 * (indicator - att_at_t3[k]);
result[k] += local_derivative * datt_t2;
}
}
for(int k = 0; k < UNROLL; ++k) {
result[k] = cg::reduce(warp, result[k], cg::plus<float>());
}
if (warp.thread_rank() < UNROLL) {
dpreatt_bth[t3 + warp.thread_rank()] = scale * result[warp.thread_rank()];
}
}
}
__global__ void softmax_autoregressive_backward_kernel5(float* __restrict__ dpreatt, const float* __restrict__ datt,
const float* __restrict__ att,
int B, int T, int C, int NH) {
constexpr int UNROLL = 8;
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
int t3 = UNROLL * (blockIdx.x * warp.meta_group_size() + warp.meta_group_rank());
int idx = blockIdx.y * T * T;
if (t3 >= T) { return; }
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
for (int t = t3; t < T; t++) {
float result[UNROLL] = {};
const float* att_bth = att + idx + t * T;
const float* datt_bth = datt + idx + t * T;
float* dpreatt_bth = dpreatt + idx + t * T;
float att_at_t3[UNROLL];
for(int k = 0; k < UNROLL; ++k) {
// if t < t3+k, we're out of bounds.
// in that case, we don't care what we read, because later on,
// we won't write the corresponding result. So just clip to
// make sure this is a valid (in-bounds) memory access.
att_at_t3[k] = att_bth[min(t, t3 + k)];
}
// the code below is actually just a for loop; except,
// we have to do something special in one iteration in
// the middle, and an if turned out to have significant
// performance impact.
// so we split the loop in three parts. Ugly, but effective.
// the beginning/end loop does the same thing, so we write the code
// just once in a lambda. In this step, we're guaranteed that
// indicator == 0
auto loop_step = [&](int t2){
float p = att_bth[t2] * datt_bth[t2];
for (int k = 0; k < UNROLL; ++k) {
result[k] -= p * att_at_t3[k];
}
};
// Now the actual loop.
{
// declare the loop iterator. Needs to be kept across the
// three different parts, so it's not a local variable in
// the for loop.
int t2 = warp.thread_rank();
// first part, as long as t2 < t3, indicator == 0
for (; t2 < t3; t2 += warp.size()) {
loop_step(t2);
}
// because k <= warp.size() (==32), the event that t3+k == t2
// has to happen at this particular step.
static_assert(UNROLL <= 32, "UNROLL is too large, this won't produce correct results.");
if (t2 <= t) {
float att_t2 = att_bth[t2];
float datt_t2 = datt_bth[t2];
float p = att_t2 * datt_t2;
for (int k = 0; k < UNROLL; ++k) {
float indicator = t2 == (t3 + k) ? 1.0f : 0.0f;
result[k] += p * (indicator - att_at_t3[k]);
}
t2 += warp.size();
}
// rest of the loop, indicator == 0 again
for (; t2 <= t; t2 += warp.size()) {
loop_step(t2);
}
}
for(int k = 0; k < UNROLL; ++k) {
result[k] = cg::reduce(warp, result[k], cg::plus<float>());
}
// when storing, we need to check that this is actually a valid result.
// here, warp.thread_rank() corresponds to `k` in the previous loops.
if (warp.thread_rank() < UNROLL && t >= t3 + warp.thread_rank()) {
dpreatt_bth[t3 + warp.thread_rank()] = scale * result[warp.thread_rank()];
}
}
}
// I want `BlockSize` to be statically known to the compiler, thus we get a template here.
// This kernel takes a step back, and looks at the original CPU code again. We have some simple outer loops
// That are independent, (b, t, h), and then the inner loops over (t2, t3) where we're combining elements -- this is
// where we can reuse data and be more efficient
// => handle b, t, h through block indices; each block does all the work for the (t2, t3) loop cooperatively.
// Now we have two nested loops, and in the inner instruction, we combine indexing from both => this calls for
// loop tiling, and lifting some of the memory ops out of the loop.
// We're in luck here; if we tile so that t3 is the outer loop, we can get a sinlge write op per result, AND also cache
// the t2-indexed part of the computation, which is the problematic one because it contains a multiplication that now we
// do not have to repeat over and over.
// => do an outer t3 loop where each thread gets one t3 index. Then, do an outer t2 loop in steps of BlockSize, and
// prepare BlockSize many elements for the inner loop. Here, each thread calculates one element and stores it in shmem.
// Then, in the inner t2 loop, each thread reads *all* the elements previously stored and does its computations.
// This way, we do 3*BlockSize loads, but BlockSize^2 computation steps => This kernel is now entirely compute bound.
// To fix up the compute issues, as above, we replace ifs in memory reading with min, and also split the inner loop
// into a large region where we don't have to calculate the indicator, and a small, costly region where we do.
template<int BlockSize>
__global__ void __launch_bounds__(BlockSize) softmax_autoregressive_backward_kernel6(float* dpreatt, const float* datt, const float* att,
int B, int T, int C, int NH) {
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
__shared__ float att_bth_s[BlockSize];
int idx = blockIdx.y;
int t = blockIdx.x;
att += idx * T * T;
datt += idx * T * T;
dpreatt += idx * T * T;
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
const float* att_bth = att + t * T;
const float* datt_bth = datt + t * T;
float* dpreatt_bth = dpreatt + t * T;
int block_steps = ceil_div(t+1, BlockSize);
// very important: This loop condition needs to be the same for all threads.
// even if a thread later on is not going to do any work, it needs to participate in the
// data loading process!
for (int t3f = 0; t3f < block_steps; ++t3f) {
int t3 = t3f * BlockSize + block.thread_rank();
float acc = 0.f;
float at3 = att_bth[t3];
for (int t2b = 0; t2b <= t; t2b += BlockSize) {
int end = min(t + 1 - t2b, BlockSize);
block.sync();
{
int t2i = block.thread_rank();
int t2 = min(t, t2b + t2i);
att_bth_s[t2i] = att_bth[t2] * datt_bth[t2];
}
block.sync();
if(t3f * BlockSize == t2b) {
for (int t2i = 0; t2i < end; t2i++) {
int t2 = t2b + t2i;
float indicator = t2 == t3 ? 1.0f : 0.0f;
acc += att_bth_s[t2i] * (indicator - at3);
}
} else {
for (int t2i = 0; t2i < end; t2i++) {
acc += att_bth_s[t2i] * (0.f - at3);
}
}
}
dpreatt_bth[t3] = scale * acc;
}
}
// Actually disentangling the loops and simplifying the resulting math gives us this pretty nice kernel.
template<int BlockSize>
__global__ void softmax_autoregressive_backward_kernel7(float* dpreatt, const float* datt, const float* att,
int B, int T, int C, float scale) {
namespace cg = cooperative_groups;
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;
int t = blockIdx.x;
att += idx * T * T;
datt += idx * T * T;
dpreatt += idx * T * T;
const float* att_bth = att + t * T;
const float* datt_bth = datt + t * T;
float* dpreatt_bth = dpreatt + t * T;
if(warp.meta_group_rank() == 0) {
block_acc[warp.thread_rank()] = 0;
}
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<float>{});
block.sync();
local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
for (int t3 = block.thread_rank(); t3 <= t; t3 += BlockSize) {
float acc = att_bth[t3] * (datt_bth[t3] - local_sum);
dpreatt_bth[t3] = scale * acc;
}
}
// The slightly less pretty version of kernel 7. Adding in all the dirty tricks that can give us a few more percent
// - streaming memory access instructions
// - reordering blocks to prevent tail effect
// - multiple values of T per block
template<int BlockSize>
__global__ void softmax_autoregressive_backward_kernel8(float* dpreatt, const float* datt, const float* att,
int B, int T, int C, float scale) {
namespace cg = cooperative_groups;
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<float>{});
block.sync();
local_sum = cg::reduce(warp, block_acc[warp.thread_rank()], cg::plus<float>{});
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);
}
}
}
// ----------------------------------------------------------------------------
// kernel launchers
// attention forward pass kernel
void attention_forward(float* out, float* vaccum, float* qkvr, float* preatt, float* att,
const float* inp,
int B, int T, int C, int NH,
const int block_size) {
// 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<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
// batched matrix multiply with cuBLAS
const float alpha = 1.0f;
const float beta = 0.0f;
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 softmax_block_size = 256;
int grid_size = ceil_div(B * NH * T * 32, softmax_block_size);
softmax_forward_kernel5<<<grid_size, softmax_block_size>>>(att, scale, preatt, B * NH, T);
// new approach: first cuBLAS another batched matmul
// vaccum = 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<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
}
void launch_softmax_1(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
int num_blocks = ceil_div(T, block_size);
softmax_autoregressive_backward_kernel1<<<dim3(num_blocks, B*NH), block_size>>>(dpreatt, datt, att, B, T, C, NH);
}
void launch_softmax_2(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
int num_blocks = ceil_div(T, block_size);
softmax_autoregressive_backward_kernel2<<<dim3(num_blocks, B*NH), block_size>>>(dpreatt, datt, att, B, T, C, NH);
}
void launch_softmax_3(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
int num_blocks = ceil_div(32*T, block_size);
softmax_autoregressive_backward_kernel3<<<dim3(num_blocks, B*NH), block_size>>>(dpreatt, datt, att, B, T, C, NH);
}
void launch_softmax_4(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
int num_blocks = ceil_div(32/8*T, block_size);
softmax_autoregressive_backward_kernel4<<<dim3(num_blocks, B*NH), block_size>>>(dpreatt, datt, att, B, T, C, NH);
}
void launch_softmax_5(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
int num_blocks = ceil_div(32/8*T, block_size);
softmax_autoregressive_backward_kernel5<<<dim3(num_blocks, B*NH), block_size>>>(dpreatt, datt, att, B, T, C, NH);
}
template<class Launcher>
void dispatch_launch(Launcher&& launch, int block_size) {
switch(block_size) {
case 32:
return launch(std::integral_constant<int, 32>{});
case 64:
return launch(std::integral_constant<int, 64>{});
case 128:
return launch(std::integral_constant<int, 128>{});
case 256:
return launch(std::integral_constant<int, 256>{});
case 512:
return launch(std::integral_constant<int, 512>{});
case 1024:
return launch(std::integral_constant<int, 1024>{});
default:
assert(false && "Invalid block size");
}
}
void launch_softmax_6(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
auto launch = [&](auto int_const) {
softmax_autoregressive_backward_kernel6<int_const.value><<<dim3(T, B * NH), int_const.value>>>(dpreatt, datt, att, B, T, C, NH);
};
dispatch_launch(launch, block_size);
}
void launch_softmax_7(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
auto launch = [&](auto int_const) {
constexpr int block_size = int_const.value;
softmax_autoregressive_backward_kernel7<block_size><<<dim3(T, B * NH), block_size>>>
(dpreatt, datt, att, B, T, C, scale);
};
dispatch_launch(launch, block_size);
}
void launch_softmax_8(float* dpreatt, float* datt, const float* att, int B, int T, int C, int NH, int block_size) {
int hs = C / NH; // head size
float scale = 1.0f / sqrtf(hs);
auto launch = [&](auto int_const) {
constexpr int block_size = int_const.value;
softmax_autoregressive_backward_kernel8<block_size><<<dim3(T / 4, B * NH), block_size>>>
(dpreatt, datt, att, B, T, C, scale);
};
dispatch_launch(launch, block_size);
}
// 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)
template<class SoftmaxKernel>
void attention_backward1(float* dinp, float* dqkvr, float* dpreatt, float* datt, float* dvaccum,
const float* dout,
const float* inp, const float* qkvr, const float* preatt, const float* att, const float* vaccum,
int B, int T, int C, int NH,
SoftmaxKernel softmax_autoregressive_backward,
const int block_size) {
int HS = C / NH; // head size
const float alpha = 1.0f;
const float beta = 1.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<<<num_blocks, block_size>>>(dvaccum, dout, B, T, NH, HS);
cudaCheck(cudaGetLastError());
// backward into datt
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
CUBLAS_OP_T, CUBLAS_OP_N,
T, T, HS,
&alpha,
v, HS, T * HS,
dvaccum, HS, T * HS,
&beta,
datt, T, T * T,
B * NH));
// backward into dv
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_T,
HS, T, T,
&alpha,
dvaccum, HS, T * HS,
att, T, T * T,
&beta,
dv, HS, T * HS,
B * NH));
// backward into preatt
softmax_autoregressive_backward(dpreatt, datt, att, B, T, C, NH, block_size);
cudaCheck(cudaGetLastError());
// backward into q
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_N,
HS, T, T,
&alpha,
k, HS, T * HS,
dpreatt, T, T * T,
&beta,
dq, HS, T * HS,
B * NH));
// backward into k
cublasCheck(cublasSgemmStridedBatched(cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_T,
HS, T, T,
&alpha,
q, HS, T * HS,
dpreatt, T, T * T,
&beta,
dk, HS, T * HS,
B * NH));
// backward into inp
num_blocks = ceil_div(B * NH * T * HS, block_size);
permute_kernel_backward<<<num_blocks, block_size>>>(dinp, dq, dk, dv, B, T, NH, HS);
cudaCheck(cudaGetLastError());
}
// kernel version dispatch
void attention_backward(int kernel_num,
float* dinp, float* dqkvr, float* dpreatt, float* datt, float* dvaccum,
const float* dout,
const float* inp, const float* qkvr, const float* preatt, const float* att, const float* vaccum,
int B, int T, int C, int NH,
const int block_size) {
switch (kernel_num) {
case 1:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_1, block_size);
break;
case 2:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_2, block_size);
break;
case 3:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_3, block_size);
break;
case 4:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_4, block_size);
break;
case 5:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_5, block_size);
break;
case 6:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_6, block_size);
break;
case 7:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_7, block_size);
break;
case 8:
attention_backward1(dinp, dqkvr, dpreatt, datt, dvaccum, dout, inp, qkvr, preatt, att, vaccum, B, T, C, NH,
launch_softmax_8, block_size);
break;
default:
printf("Invalid kernel number\n");
exit(1);
}
}
// ----------------------------------------------------------------------------
int main(int argc, char **argv) {
setup_main();
// hyperparameters
int B = 4;
int T = 1024;
int C = 768;
int NH = 12;
// read kernel_num from command line
int kernel_num = 1;
if (argc > 1) {
kernel_num = atoi(argv[1]);
}
printf("Using kernel %d\n", kernel_num);
// create the host memory for the forward pass
float* inp = make_random_float(B * T * 3 * C);
float* qkvr = (float*)malloc(B * T * 3 * C * sizeof(float));
float* preatt = (float*)malloc(B * NH * T * T * sizeof(float));
float* att = (float*)malloc(B * NH * T * T * sizeof(float));
float* vaccum = (float*)malloc(B * T * C * sizeof(float));
float* out = (float*)malloc(B * T * C * sizeof(float));
// execute the forward pass on the CPU
attention_forward_cpu(out, preatt, att, inp, B, T, C, NH);
// create device memory for the forward pass
float *d_inp, *d_qkvr, *d_preatt, *d_att, *d_vaccum, *d_out;
cudaCheck(cudaMalloc(&d_inp, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_qkvr, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_preatt, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMalloc(&d_att, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMalloc(&d_vaccum, B * T * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float)));
// copy over the input
cudaCheck(cudaMemcpy(d_inp, inp, B * T * 3 * C * sizeof(float), cudaMemcpyHostToDevice));
// execute the forward pass on the GPU
const int block_size = 256;
attention_forward(d_out, d_vaccum, d_qkvr, d_preatt, d_att, d_inp, B, T, C, NH, block_size);
// check that preatt, att, and out match between the CPU and GPU versions
printf("Checking the forward pass CPU <-> GPU...\n");
printf("[preatt]\n"); validate_result(d_preatt, preatt, "preatt", B * T * C, 5e-3f);
printf("[att]\n"); validate_result(d_att, att, "att", B * T * C, 1e-3f);
printf("[out]\n"); validate_result(d_out, out, "out", B * T * C, 1e-3f);
// set up the memory for the backward pass
float* dout = make_random_float(B * T * C); // the gradients on the output
float* dinp = make_zeros_float(B * T * 3 * C); // zeros for all else, to += into
float* dpreatt = make_zeros_float(B * NH * T * T);
float* datt = make_zeros_float(B * NH * T * T);
// call backward() on the CPU to get our reference gradients
attention_backward_cpu(dinp, dpreatt, datt, dout, inp, att, B, T, C, NH);
// create device memory for the backward pass
float *d_dinp, *d_dqkvr, *d_dpreatt, *d_datt, *d_dvaccum, *d_dout;
cudaCheck(cudaMalloc(&d_dinp, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_dqkvr, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_dpreatt, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMalloc(&d_datt, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMalloc(&d_dvaccum, B * T * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_dout, B * T * C * sizeof(float)));
// copy over the dout gradients that starts the backprop chain
cudaCheck(cudaMemcpy(d_dout, dout, B * T * C * sizeof(float), cudaMemcpyHostToDevice));
// memset all the other memory to zeros, to += into
cudaCheck(cudaMemset(d_dinp, 0, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMemset(d_dqkvr, 0, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMemset(d_dpreatt, 0, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMemset(d_datt, 0, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMemset(d_dvaccum, 0, B * T * C * sizeof(float)));
// call backward() on the GPU
attention_backward(kernel_num, d_dinp, d_dqkvr, d_dpreatt, d_datt, d_dvaccum,
d_dout, d_inp, d_qkvr, d_preatt, d_att, d_vaccum,
B, T, C, NH, block_size);
// check that the gradients match between the CPU and GPU versions
// note that we will only check the correctness at [att, preatt, inp]
// the gradients at qkvr and vaccum will remain unchecked, but are
// assumed to be correct if the other gradients are correct
printf("Checking the backward pass CPU <-> GPU...\n");
printf("[datt]\n"); validate_result(d_datt, datt, "datt", B * NH * T * T, 5e-3f);
printf("[dpreatt]\n"); validate_result(d_dpreatt, dpreatt, "dpreatt", B * NH * T * T, 1e-3f);
printf("[dinp]\n"); validate_result(d_dinp, dinp, "dinp", B * T * 3 * C, 1e-3f);
// also let's manually step through the gradients here
float* h_dinp = (float*)malloc(B * T * 3 * C * sizeof(float));
cudaCheck(cudaMemcpy(h_dinp, d_dinp, B * T * 3 * C * sizeof(float), cudaMemcpyDeviceToHost));
int num_match = 0;
int num_no_match = 0;
int num_zero_grad = 0;
int HS = C / NH;
for (int i = 0; i < B * T * 3 * C; i++) {
// the dimensions of inp are (B, T, 3, NH, HS)
// where B = batch, T = time, 3 = qkv, NH = num heads, HS = head size
// unpack the individual b,t,qkvix,h,c indices
int ix = i;
int c = ix % HS;
ix /= HS;
int h = ix % NH;
ix /= NH;
int qkvix = ix % 3;
ix /= 3;
int t = ix % T;
ix /= T;
int b = ix;
float diff = fabs(dinp[i] - h_dinp[i]);
// attempt to index at random
if (b == 1 && t == 5 && c == 23 && h == 2) {
printf("ix %5d [b=%4d, t=%4d, qkv=%4d, nh=%4d, hs=%4d]: ref: %f gpu: %f\n", i, b, t, qkvix, h, c, dinp[i], h_dinp[i]);
}
if (diff > 1e-4f) {
num_no_match++;
} else {
num_match++;
}
if (dinp[i] == 0.0f) {
num_zero_grad++;
}
}
printf("Number of matching gradients: %d (%.2f%% of total)\n", num_match, 100*(float)num_match / (B * T * 3 * C));
printf("Number of non-matching gradients: %d (%.2f%% of total)\n", num_no_match, 100*(float)num_no_match / (B * T * 3 * C));
printf("Number of gradients that are exactly zero: %d (%.2f%% of total)\n", num_zero_grad, 100*(float)num_zero_grad / (B * T * 3 * C));
// final verdict
printf("All results match. Starting benchmarks.\n\n");
// benchmark speed of the kernel
int block_sizes[] = {32, 64, 128, 256, 512, 1024};
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
int block_size = block_sizes[j];
int repeat_times = 10;
float elapsed_time = benchmark_kernel(repeat_times, attention_backward,
kernel_num, d_dinp, d_dqkvr, d_dpreatt, d_datt, d_dvaccum,
d_dout, d_inp, d_qkvr, d_preatt, d_att, d_vaccum,
B, T, C, NH, block_size);
printf("block_size %4d | time %f ms\n", block_size, elapsed_time);
}
// free memory
free(inp);
free(qkvr);
free(preatt);
free(att);
free(vaccum);
free(out);
free(dout);
free(dinp);
free(dpreatt);
free(datt);
free(h_dinp);
cudaCheck(cudaFree(d_inp));
cudaCheck(cudaFree(d_qkvr));
cudaCheck(cudaFree(d_preatt));
cudaCheck(cudaFree(d_att));
cudaCheck(cudaFree(d_vaccum));
cudaCheck(cudaFree(d_out));
cudaCheck(cudaFree(d_dinp));
cudaCheck(cudaFree(d_dqkvr));
cudaCheck(cudaFree(d_dpreatt));
cudaCheck(cudaFree(d_datt));
cudaCheck(cudaFree(d_dvaccum));
cudaCheck(cudaFree(d_dout));
cublasDestroy(cublas_handle);
return 0;
}
================================================
FILE: dev/cuda/attention_forward.cu
================================================
/*
Kernels for attention forward pass.
If you do not have CUDNN, you can remove ENABLE_CUDNN to run the other kernels
See the README for cuDNN install instructions
Compile example with cuDNN:
nvcc -I/PATH/TO/cudnn-frontend/include -DENABLE_CUDNN -O3 --use_fast_math --lcublas -lcublasLt -lcudnn attention_forward.cu -o attention_forward
Compile example without cuDNN:
nvcc -O3 --use_fast_math -lcublas -lcublasLt attention_forward.cu -o attention_forward
version 1 is naive port from CPU code to kernel, parallelize over batch, time, heads only
./attention_forward 1
version 2 is a naive implementation of flash attention, taken, adapted from
https://github.com/tspeterkim/flash-attention-minimal
and with help from
https://github.com/leloykun/flash-hyperbolic-attention-minimal
sadly, this flash attention version seems about 3X slower than the naive version
./attention_forward 2
version 3 is a cuBLAS + softmax version, similar to the PyTorch implementation
cuBLAS is used both to calculate the QK^T and the final weighted sum
the softmax is calculated using a custom, efficient kernel as well
this turns out to be ~20X faster than (1) nice
./attention_forward 3
version 4 is a further optimized kernel that fuses the scale operation,
uses a directly autoregressive softmax, and uses the online softmax algorithm.
./attention_forward 4
version 5 is a FP16 version of kernel 4
./attention_forward 5
version 6 is kernel 5 skipping (un)permute (unrealistic but useful comparison point)
version 10 is using cuDNN Flash Attention using FP16 or BF16, see:
https://github.com/NVIDIA/cudnn-frontend/blob/main/docs/operations/Attention.md
./attention_forward 10
version 11 is kernel 10 skipping FP16/FP32 conversions (full FP16/BF16 network)
./attention_forward 11
*/
//#define ENABLE_CUDNN // can be enabled via nvcc "-DENABLE_CUDNN"
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <float.h>
#include <cublas_v2.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#define ENABLE_BF16
#include "common.h"
// ----------------------------------------------------------------------------
// CUDA & cuDNN setup
static bool first_run_validation = true; // always run e.g. permute on 1st run
#ifdef ENABLE_CUDNN
#include <cudnn_frontend.h>
namespace fe = cudnn_frontend;
#if CUBLAS_LOWP == CUDA_R_16BF
#define CUDNN_16BIT fe::DataType_t::BFLOAT16
#else
#define CUDNN_16BIT fe::DataType_t::HALF
#endif
static cudnnHandle_t cudnn_handle;
static size_t cudnn_workspace_size = 0; // dynamically allocated as needed (up to 256MiB!)
static void* cudnn_workspace = NULL;
#define checkCudaErr(err) assert((int)err == 0);
#define checkCudnnErr(err) assert((int)err == 0);
#endif // ENABLE_CUDNN
// ----------------------------------------------------------------------------
// CPU code reference
void attention_forward_cpu(float* out, float* preatt, float* att,
const float* inp,
int B, int T, int C, int NH) {
// input is (B, T, 3C) Q,K,V
// preatt, att are (B, NH, T, T)
// output is (B, T, C)
int C3 = C*3;
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* preatt_bth = preatt + b*NH*T*T + h*T*T + t*T;
float* att_bth = att + b*NH*T*T + h*T*T + t*T;
// pass 1: calculate query dot key and maxval
float maxval = -FLT_MAX;
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
// (query_t) dot (key_t2)
float val = 0.0f;
for (int i = 0; i < hs; i++) {
val += query_t[i] * key_t2[i];
}
val *= scale;
if (val > maxval) {
maxval = val;
}
preatt_bth[t2] = val;
}
// pad with -INFINITY outside of autoregressive region for debugging comparisons
for (int t2 = t+1; t2 < T; t2++) {
preatt_bth[t2] = -INFINITY;
}
// pass 2: calculate the exp and keep track of sum
float expsum = 0.0f;
for (int t2 = 0; t2 <= t; t2++) {
float expv = expf(preatt_bth[t2] - maxval);
expsum += expv;
att_bth[t2] = expv;
}
float expsum_inv = expsum == 0.0f ? 0.0f : 1.0f / expsum;
// pass 3: normalize to get the softmax
for (int t2 = 0; t2 < T; t2++) {
if (t2 <= t) {
att_bth[t2] *= expsum_inv;
} else {
// causal attention mask. not strictly necessary to set to zero here
// only doing this explicitly for debugging and checking to PyTorch
att_bth[t2] = 0.0f;
}
}
// pass 4: accumulate weighted values into the output of attention
float* out_bth = out + b * T * C + t * C + h * hs;
for (int i = 0; i < hs; i++) { out_bth[i] = 0.0f; }
for (int t2 = 0; t2 <= t; t2++) {
const float* value_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C*2; // +C*2 because it's value
float att_btht2 = att_bth[t2];
for (int i = 0; i < hs; i++) {
out_bth[i] += att_btht2 * value_t2[i];
}
}
}
}
}
}
// ----------------------------------------------------------------------------
// GPU kernels
__global__ void attention_query_key_kernel1(float* preatt, const float* inp,
int B, int T, int C, int NH) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total_threads = B * NH * T * T;
if (idx < total_threads) {
int t2 = idx % T;
int t = (idx / T) % T;
if (t2 > t) {
// autoregressive mask
preatt[idx] = -INFINITY;
return;
}
int h = (idx / (T * T)) % NH;
int b = idx / (NH * T * T);
int C3 = C*3;
int hs = C / NH; // head size
const float* query_t = inp + b * T * C3 + t * C3 + h * hs;
const float* key_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C; // +C because it's key
// (query_t) dot (key_t2)
float val = 0.0f;
for (int i = 0; i < hs; i++) {
val += query_t[i] * key_t2[i];
}
val *= 1.0 / sqrtf(hs);
preatt[idx] = val;
}
}
__global__ void attention_softmax_kernel1(float* att, const float* preatt,
int B, int T, int NH) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total_threads = B * T * NH;
if (idx < total_threads) {
int h = idx % NH;
int t = (idx / NH) % T;
int b = idx / (NH * T);
const float* preatt_bth = preatt + b*NH*T*T + h*T*T + t*T;
float* att_bth = att + b*NH*T*T + h*T*T + t*T;
// find maxval
float maxval = -FLT_MAX;
for (int t2 = 0; t2 <= t; t2++) {
if (preatt_bth[t2] > maxval) {
maxval = preatt_bth[t2];
}
}
// calculate the exp and keep track of sum
float expsum = 0.0f;
for (int t2 = 0; t2 <= t; t2++) {
float expv = expf(preatt_bth[t2] - maxval);
expsum += expv;
att_bth[t2] = expv;
}
float expsum_inv = expsum == 0.0f ? 0.0f : 1.0f / expsum;
// normalize to get the softmax
for (int t2 = 0; t2 < T; t2++) {
if (t2 <= t) {
att_bth[t2] *= expsum_inv;
} else {
// causal attention mask. not strictly necessary to set to zero here
// only doing this explicitly for debugging and checking to PyTorch
att_bth[t2] = 0.0f;
}
}
}
}
// warp-level reduction for finding the maximum value
__device__ float warpReduceMax(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val = fmaxf(val, __shfl_down_sync(0xFFFFFFFF, val, offset));
}
return val;
}
__global__ void softmax_forward_kernel4(float* out, const float* inp, int N, int C) {
// out is (N, C) just like inp. Each row of inp will get softmaxed.
// same as kernel3, but can handle any block size (multiple of 32)
// each row of C elements is handled by block_size threads
// furthermore, each block_size threads get executed in warps of 32 threads
// special reduction operations warpReduceMax/warpReduceSum are used for intra-warp reductions
// shared memory is used for inter-warp reduction
extern __shared__ float shared[];
int idx = blockIdx.x;
int tid = threadIdx.x;
int warpId = threadIdx.x / 32; // warp index within a block
int laneId = threadIdx.x % 32; // thread index within a warp
// 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];
// one row of inp, i.e. inp[idx, :] of shape (C,)
const float* x = inp + idx * C;
// first, thread coarsening by directly accessing global memory in series
float maxval = -INFINITY;
for (int i = tid; i < C; i += blockDim.x) {
maxval = fmaxf(maxval, x[i]);
}
// now within-warp reductions for maxval
maxval = warpReduceMax(maxval);
// the 0th thread of each warp writes the maxval of that warp to shared memory
if (laneId == 0) maxvals[warpId] = maxval;
__syncthreads();
// now the 0th thread reduces the maxvals in shared memory, i.e. across warps
if (tid == 0) {
float val = maxvals[tid];
for (int i = 1; i < warpsPerBlock; i++) {
val = fmaxf(val, maxvals[i]);
}
// store the final max in the first position
maxvals[0] = val;
}
__syncthreads();
// broadcast the max to all threads
float offset = maxvals[0];
// compute expf and write the result to global memory
for (int i = tid; i < C; i += blockDim.x) {
// subtract max for numerical stability
out[idx * C + i] = expf(x[i] - offset);
}
// okay now we calculated exp(x - max(x))
// step 2: sum all the values and divide by the sum
// thread coarsening for sum
x = out + idx * C;
float sumval = 0.0f;
for (int i = tid; i < C; i += blockDim.x) {
sumval += x[i];
}
// within-warp reduction for sumval
sumval = warpReduceSum(sumval);
// write sumval to shared memory
if (laneId == 0) sumvals[warpId] = sumval;
__syncthreads();
// inter-thread reduction of sum
if (tid == 0) {
float val = sumvals[tid];
for (int i = 1; i < warpsPerBlock; ++i) {
val += sumvals[i];
}
sumvals[0] = val;
}
__syncthreads();
// broadcast the sum to all threads
float sum = sumvals[0];
// divide the whole row by the sum
for (int i = tid; i < C; i += blockDim.x) {
out[idx * C + i] = x[i] / sum;
}
}
__device__ float& vec_at(float4& vec, int index) {
return reinterpret_cast<float*>(&vec)[index];
}
__device__ float vec_at(const float4& vec, int index) {
return reinterpret_cast<const float*>(&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);
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();
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<const float4*>(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<float>{});
sumval *= expf(inv_temperature * (maxval - global_maxval));
float sum = cg::reduce(warp, sumval, cg::plus<float>{});
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 attention_value_kernel1(float* out, const float* att, const float* inp,
int B, int T, int C, int NH) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total_threads = B * T * NH;
if (idx < total_threads) {
int h = idx % NH;
int t = (idx / NH) % T;
int b = idx / (NH * T);
int C3 = C*3;
int hs = C / NH; // head size
float* out_bth = out + b * T * C + t * C + h * hs;
const float* att_bth = att + b*NH*T*T + h*T*T + t*T;
for (int i = 0; i < hs; i++) { out_bth[i] = 0.0f; }
for (int t2 = 0; t2 <= t; t2++) {
const float* value_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C*2; // +C*2 because it's value
float att_btht2 = att_bth[t2];
for (int i = 0; i < hs; i++) {
out_bth[i] += att_btht2 * value_t2[i];
}
}
}
}
__global__
void attention_forward_kernel2(
const float* Q,
const float* K,
const float* V,
const int N,
const int d,
const int Tc,
const int Tr,
const int Bc,
const int Br,
const float softmax_scale,
float* l,
float* m,
float* O
) {
int tx = threadIdx.x;
int bx = blockIdx.x; int by = blockIdx.y; // batch and head index
// Offset into Q,K,V,O,l,m - different for each batch and head
int qkv_offset = (bx * gridDim.y * N * d) + (by * N * d); // gridDim.y = nh
int lm_offset = (bx * gridDim.y * N) + (by * N); // offset for l and m
// Define SRAM for Q,K,V,S
extern __shared__ float sram[];
int tile_size = Bc * d; // size of Qi, Kj, Vj
float* Qi = sram;
float* Kj = &sram[tile_size];
float* Vj = &sram[tile_size * 2];
float* S = &sram[tile_size * 3];
for (int j = 0; j < Tc; j++) {
// Load Kj, Vj to SRAM
for (int x = 0; x < d; x++) {
Kj[(tx * d) + x] = K[qkv_offset + (tile_size * j) + (tx * d) + x];
Vj[(tx * d) + x] = V[qkv_offset + (tile_size * j) + (tx * d) + x];
}
__syncthreads(); // such that the inner loop can use the correct Kj, Vj
for (int i = 0; i < Tr; i++) {
// if past the end of the sequence, break
if (i * Br + tx >= N) {
break;
}
// Load Qi to SRAM, l and m to registers
for (int x = 0; x < d; x++) {
Qi[(tx * d) + x] = Q[qkv_offset + (tile_size * i) + (tx * d) + x];
}
float row_m_prev = m[lm_offset + (Br * i) + tx];
float row_l_prev = l[lm_offset + (Br * i) + tx];
// S = QK^T, row_m = rowmax(S)
// S[tx][y] = Sum_{x = 0}^{d-1} {Qi[tx][x] * Kj[y][x]}
// row_m = Max_{y = 0}^{Bc-1} S[tx][y]
// with causal masking
float row_m = -INFINITY;
for (int y = 0; y < Bc; y++) {
if (j * Bc + y >= N) {
break;
}
float sum = 0;
for (int x = 0; x < d; x++) {
sum += Qi[(tx * d) + x] * Kj[(y * d) + x];
}
sum *= softmax_scale;
if (i * Br + tx < j * Bc + y)
sum = -INFINITY;
S[(Bc * tx) + y] = sum;
if (sum > row_m)
row_m = sum;
}
// implement softmax with causal masking
// P = exp(S - row_m), row_l = rowsum(P)
// P[tx][y] = exp(S[tx][y] - row_m)
float row_l = 0;
for (int y = 0; y < Bc; y++) {
if (j * Bc + y >= N) {
break;
}
if (i * Br + tx < j * Bc + y)
S[(Bc * tx) + y] = 0;
else
S[(Bc * tx) + y] = __expf(S[(Bc * tx) + y] - row_m);
row_l += S[(Bc * tx) + y];
}
// Compute new m and l
float row_m_new = max(row_m_prev, row_m);
float row_l_new = (__expf(row_m_prev - row_m_new) * row_l_prev) + (__expf(row_m - row_m_new) * row_l);
// Write O, l, m to HBM
for (int x = 0; x < d; x++) {
float pv = 0; // Pij * Vj
for (int y = 0; y < Bc; y++) {
if (j * Bc + y >= N) {
break;
}
pv += S[(Bc * tx) + y] * Vj[(y * d) + x];
}
O[qkv_offset + (tile_size * i) + (tx * d) + x] = (1 / row_l_new) \
* ((row_l_prev * __expf(row_m_prev - row_m_new) * O[qkv_offset + (tile_size * i) + (tx * d) + x]) \
+ (__expf(row_m - row_m_new) * pv));
}
m[lm_offset + (Br * i) + tx] = row_m_new;
l[lm_offset + (Br * i) + tx] = row_l_new;
}
__syncthreads(); // otherwise, thread can use the wrong Kj, Vj in inner loop
}
}
__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] = inp[inp_idx];
k[idx] = inp[inp_idx + NH * d];
v[idx] = inp[inp_idx + 2 * (NH * d)];
}
}
__global__ void unpermute_kernel(const 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] = inp[idx];
}
}
__global__ void scale_kernel(float* inp, float scale, int B, int NH, int T) {
// scales the pre-softmax attention scores by scale
// and sets the autoregressive locations to -INFINITY
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < B * NH * T * T) {
int rest = idx % (NH * T * T);
rest = rest % (T * T);
int t2 = rest / T;
int t = rest % T;
if (t > t2) {
inp[idx] = -INFINITY;
} else {
inp[idx] *= scale;
}
}
}
// direct translation of the CPU kernel. Each warp handles ont (b, h, t) combination.
// The important changes compared to the CPU version:
// - each inner loop is handled by a warp
// - don't write non-autoregressive parts
// - reordered the last loops so that we can do all writing in the outer loop.
__global__ void attention_forward_fused1(float* out, float* preatt, float* att,
const float* inp,
int B, int T, int C, int NH) {
// input is (B, T, 3C) Q,K,V
// preatt, att are (B, NH, T, T)
// output is (B, T, C)
int C3 = C*3;
int hs = C / NH; // head size
float scale = 1.0 / sqrtf(hs);
namespace cg = cooperative_groups;
cg::thread_block block = cg::this_thread_block();
cg::thread_block_tile<32> warp = cg::tiled_partition<32>(block);
int t = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
int h = blockIdx.y;
int b = blockIdx.z;
if(t >= T) return;
const float* query_t = inp + b * T * C3 + t * C3 + h * hs;
float* preatt_bth = preatt + b*NH*T*T + h*T*T + t*T;
float* att_bth = att + b*NH*T*T + h*T*T + t*T;
// pass 1: calculate query dot key and maxval
float maxval = -INFINITY;
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
// (query_t) dot (key_t2)
float val = 0.0f;
for (int i = warp.thread_rank(); i < hs; i += warp.size()) {
val += query_t[i] * key_t2[i];
}
val = cg::reduce(warp, val, cg::plus<float>{});
val *= scale;
maxval = max(maxval, val);
if(warp.thread_rank() == 0) {
preatt_bth[t2] = val;
}
}
// pass 2: calculate the exp and keep track of sum
float expsum = 0.0f;
for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {
float expv = expf(preatt_bth[t2] - maxval);
expsum += expv;
}
expsum = cg::reduce(warp, expsum, cg::plus<float>{});
float expsum_inv = expsum == 0.0f ? 0.0f : 1.0f / expsum;
// pass 3: normalize to get the softmax is combined with the next loop to reduce memory round-trips
for (int t2 = warp.thread_rank(); t2 <= t; t2 += warp.size()) {
att_bth[t2] = expf(preatt_bth[t2] - maxval) * expsum_inv;
}
// pass 4: accumulate weighted values into the output of attention
float* out_bth = out + b * T * C + t * C + h * hs;
for (int i = warp.thread_rank(); i < hs; i += warp.size()) {
float o = 0.f;
for (int t2 = 0; t2 <= t; t2++) {
const float* value_t2 = inp + b * T * C3 + t2 * C3 + h * hs + C * 2; // +C*2 because it's value
float att_btht2 = att_bth[t2];
o += att_btht2 * value_t2[i];
}
out_bth[i] = o;
}
}
// ----------------------------------------------------------------------------
// kernel launcher
void attention_forward1(float* out, float* preatt, float* att,
const float* inp,
int B, int T, int C, int NH,
const int block_size) {
// attention calculation
int total_threads = B * NH * T * T;
int num_blocks = ceil_div(total_threads, block_size);
attention_query_key_kernel1<<<num_blocks, block_size>>>(preatt, inp, B, T, C, NH);
// softmax and value accumulation
total_threads = B * T * NH;
num_blocks = ceil_div(total_threads, block_size);
attention_softmax_kernel1<<<num_blocks, block_size>>>(att, preatt, B, T, NH);
attention_value_kernel1<<<num_blocks, block_size>>>(out, att, inp, B, T, C, NH);
}
void attention_forward2(float* out,
const float* inp,
int B, int T, int C, int NH,
const int block_size) {
// TODO there should be no mallocs inside any of these functions!
// not fixing this because we don't intend to use attention_forward2,
// it seems to be way too slow as is
// these are hardcoded to 32 for now
const int Bc = 32;
const int Br = 32;
// renaming these to be consistent with the kernel
// const int B = B;
const int nh = NH;
const int N = T;
const int d = C / NH;
// more
const int Tc = ceil((float) N / Bc);
const int Tr = ceil((float) N / Br);
const float softmax_scale = 1.0 / sqrt(d);
// create some temporary memory
float* l;
float* m;
cudaCheck(cudaMalloc(&l, B * nh * N * sizeof(float)));
cudaCheck(cudaMalloc(&m, B * nh * N * sizeof(float)));
cudaCheck(cudaMemset(l, 0, B * nh * N * sizeof(float)));
cudaCheck(cudaMemset(m, -10000.0f, B * nh * N * sizeof(float)));
// calculate SRAM size needed per block, ensure we have enough shared memory
int col_tile_size = Bc * d; // size of Kj, Vj
int row_tile_size = Br * d; // size of Qi
const int sram_size =
(2 * col_tile_size * sizeof(float)) // SRAM size for Kj, Vj
+ (row_tile_size * sizeof(float)) // SRAM size for Qi
+ (Bc * Br * sizeof(float)); // SRAM size for S
int max_sram_size;
cudaDeviceGetAttribute(&max_sram_size, cudaDevAttrMaxSharedMemoryPerBlock, 0);
if (sram_size > max_sram_size) {
printf("Max shared memory: %d, requested shared memory: %d \n", max_sram_size, sram_size);
printf("SRAM size exceeds maximum shared memory per block\n");
printf("Try decreasing col_tile_size or row_tile_size further\n");
exit(1);
}
// grid and block dims
dim3 grid_dim(B, nh); // batch_size x num_heads
dim3 block_dim(Br); // Br threads per block
// 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)
// so we have to permute the tensor using a kernel with block_size
float *q, *k, *v;
cudaCheck(cudaMalloc(&q, B * T * C * sizeof(float)));
cudaCheck(cudaMalloc(&k, B * T * C * sizeof(float)));
cudaCheck(cudaMalloc(&v, B * T * C * sizeof(float)));
int total_threads = B * N * nh * d;
int num_blocks = ceil_div(total_threads, block_size);
permute_kernel<<<num_blocks, block_size>>>(q, k, v, inp, B, N, nh, d);
// now actually call the flash attention kernel
attention_forward_kernel2<<<grid_dim, block_dim, sram_size>>>(
q, k, v,
N, d, Tc, Tr, Bc, Br, softmax_scale,
l, m, out
);
// out has shape (B, nh, N, d) but we need to unpermute it to (B, N, nh, d)
unpermute_kernel<<<num_blocks, block_size>>>(out, q, B, N, nh, d);
cudaCheck(cudaMemcpy(out, q, B * T * C * sizeof(float), cudaMemcpyDeviceToDevice));
// free memory
cudaCheck(cudaFree(l));
cudaCheck(cudaFree(m));
cudaCheck(cudaFree(q));
cudaCheck(cudaFree(k));
cudaCheck(cudaFree(v));
}
void attention_forward3(float* out, float* vaccum, float* qkvr, float* preatt, float* att,
const float* inp,
int B, int T, int C, int NH,
const int block_size) {
// 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<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
// batched matrix multiply with cuBLAS
const float alpha = 1.0f;
const float beta = 0.0f;
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.0f / sqrtf(HS);
total_threads = B * NH * T * T;
num_blocks = ceil_div(total_threads, block_size);
scale_kernel<<<num_blocks, block_size>>>(preatt, scale, B, NH, T);
// softmax. preatt is (B, NH, T, T) but we view it as (B * NH * T, T) and use the softmax kernel
int softmax_block_size = 256;
int grid_size = B * NH * T;
size_t shared_mem_size = 2 * softmax_block_size / 32 * sizeof(float);
softmax_forward_kernel4<<<grid_size, softmax_block_size, shared_mem_size>>>(att, preatt, B * NH * T, T);
// new approach: first cuBLAS another batched matmul
// 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<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
}
void attention_forward4(float* out, float* vaccum, float* qkvr, float* preatt, float* att,
const float* inp,
int B, int T, int C, int NH,
const int block_size) {
// 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<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
// batched matrix multiply with cuBLAS
const float alpha = 1.0f;
const float beta = 0.0f;
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 softmax_block_size = 256;
int grid_size = ceil_div(B * NH * T * 32, softmax_block_size);
softmax_forward_kernel5<<<grid_size, softmax_block_size>>>(att, scale, preatt, B * NH, T);
// new approach: first cuBLAS another batched matmul
// 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<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
}
__global__ void softmax_forward_kernel5_lowp(floatX* out, float inv_temperature,
const floatX* 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);
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();
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 floatX* 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;
// Same thing but without float4, one at a time
for (int i = warp.thread_rank(); i < pos_by_4; i += warp.size()) {
float old_maxval = maxval;
for(int k = 0; k < 4; ++k) {
maxval = fmaxf(maxval, (float)x[4*i + k]);
}
sumval *= expf(inv_temperature * (old_maxval - maxval));
for(int k = 0; k < 4; ++k) {
sumval += expf(inv_temperature * ((float)x[4*i + k] - maxval));
}
}
if(4*pos_by_4 + warp.thread_rank() <= own_pos) {
float old_maxval = maxval;
maxval = fmaxf(maxval, (float)x[4*pos_by_4 + warp.thread_rank()]);
sumval *= expf(inv_temperature * (old_maxval - maxval));
sumval += expf(inv_temperature * ((float)x[4*pos_by_4 + warp.thread_rank()] - maxval));
}
float global_maxval = cg::reduce(warp, maxval, cg::greater<float>{});
sumval *= expf(inv_temperature * (maxval - global_maxval));
float sum = cg::reduce(warp, sumval, cg::plus<float>{});
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 * ((float)__ldcs(x + i) - global_maxval));
__stcs(out + idx * T + i, (floatX)(ev * norm));
}
}
__global__ void permute_kernel_lowp(floatX* q, floatX* k, floatX* 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] = (floatX)inp[inp_idx];
k[idx] = (floatX)inp[inp_idx + NH * d];
v[idx] = (floatX)inp[inp_idx + 2 * (NH * d)];
}
}
__global__ void unpermute_kernel_lowp(const floatX* 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] = (float)inp[idx];
}
}
void attention_forward5(float* out, floatX* vaccum, floatX* qkvr, floatX* preatt, floatX* att,
const float* inp,
int B, int T, int C, int NH,
const int block_size, bool skip_permute=false) {
// FP16 version of kernel 4 (with permute/unpermute doing FP32<->FP16)
// That permute can be skipped on perf runs to analyse its performance impact
// inp is (B, T, 3C) QKV
// preatt, att are (B, NH, T, T)
// output is (B, T, C)
// permute and separate inp from (B, T, 3, NH, HS) to 3X (B, NH, T, HS)
int HS = C / NH; // head size
floatX *q = qkvr + 0 * B * T * C;
floatX *k = qkvr + 1 * B * T * C;
floatX* v = qkvr + 2 * B * T * C;
int total_threads = B * NH * T * HS;
int num_blocks = ceil_div(total_threads, block_size);
if (!skip_permute || first_run_validation) {
permute_kernel_lowp<<<num_blocks, block_size>>>(q, k, v, inp, B, T, NH, HS);
}
// IMPORTANT: alpha/beta are FP32 for CUBLAS_COMPUTE_32F even if FP16 inputs/outputs
// But need FP16 scale for CUBLAS_COMPUTE_16F (no errors otherwise, just garbage results *sigh*)
const float alpha = 1.0f;
const float beta = 0.0f;
const floatX alpha_lowp = (floatX)alpha;
const floatX beta_lowp = (floatX)beta;
void* alpha_ptr = CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F ? (void*)&alpha_lowp : (void*)α
void* beta_ptr = CUBLAS_LOWP_COMPUTE == CUBLAS_COMPUTE_16F ? (void*)&beta_lowp : (void*)β
// batched matrix multiply with cuBLAS
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
CUBLAS_OP_T, CUBLAS_OP_N,
T, T, HS,
alpha_ptr,
k, CUBLAS_LOWP, HS, T * HS,
q, CUBLAS_LOWP, HS, T * HS,
beta_ptr,
preatt, CUBLAS_LOWP, T, T * T,
B * NH,
CUBLAS_LOWP_COMPUTE,
CUBLAS_GEMM_DEFAULT));
// multiply all elements of preatt elementwise by scale
float scale = 1.0f / sqrtf(HS);
int softmax_block_size = 256;
int grid_size = ceil_div(B * NH * T * 32, softmax_block_size);
softmax_forward_kernel5_lowp<<<grid_size, softmax_block_size>>>(att, scale, preatt, B * NH, T);
// new approach: first cuBLAS another batched matmul
// y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)
cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
CUBLAS_OP_N, CUBLAS_OP_N,
HS, T, T,
alpha_ptr,
v, CUBLAS_LOWP, HS, T * HS,
att, CUBLAS_LOWP, T, T * T,
beta_ptr,
vaccum, CUBLAS_LOWP, HS, T * HS,
B * NH,
CUBLAS_LOWP_COMPUTE,
CUBLAS_GEMM_DEFAULT));
// 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);
if(!skip_permute || first_run_validation) {
unpermute_kernel_lowp<<<num_blocks, block_size>>>(vaccum, out, B, T, NH, HS);
}
}
#ifdef ENABLE_CUDNN
using graph_tensors_fwd = std::tuple<std::shared_ptr<fe::graph::Graph>,
std::shared_ptr<fe::graph::Tensor_attributes>, // Q,
std::shared_ptr<fe::graph::Tensor_attributes>, // K,
std::shared_ptr<fe::graph::Tensor_attributes>, // V,
std::shared_ptr<fe::graph::Tensor_attributes>, // Attn_scale,
std::shared_ptr<fe::graph::Tensor_attributes>, // O
std::shared_ptr<fe::graph::Tensor_attributes>>; // Stats
// Need a cache because graph->build_operation_graph() is slow but everything else seems fast
using cache_type_fwd = std::unordered_map<std::size_t, graph_tensors_fwd>;
// Loosely based on cuDNN frontend samples functions and massively simplified
template <typename... Args>
auto lookup_cache_or_build_graph_fwd(Args... args) {
static cache_type_fwd user_maintained_cache_fwd;
auto [B, H, T, HS, is_inference_only] = std::make_tuple(args...);
auto graph = std::make_shared<fe::graph::Graph>();
graph->set_io_data_type(CUDNN_16BIT)
.set_intermediate_data_type(fe::DataType_t::FLOAT)
.set_compute_data_type(fe::DataType_t::FLOAT);
// QKV is (B, T, 3, NH, HS) which cuDNN can handle directly without an external permute
auto Q = graph->tensor(fe::graph::Tensor_attributes()
.set_name("Q")
.set_dim({B, H, T, HS})
.set_stride({3 * H * HS * T, HS, 3 * H * HS, 1}));
auto K = graph->tensor(fe::graph::Tensor_attributes()
.set_name("K")
.set_dim({B, H, T, HS})
.set_stride({3 * H * HS * T, HS, 3 * H * HS, 1}));
auto V = graph->tensor(fe::graph::Tensor_attributes()
.set_name("V")
.set_dim({B, H, T, HS})
.set_stride({3 * H * HS * T, HS, 3 * H * HS, 1}));
auto attn_scale = graph->tensor(fe::graph::Tensor_attributes()
.set_name("attn_scale")
.set_dim({1, 1, 1, 1})
.set_stride({1, 1, 1, 1})
.set_is_pass_by_value(true)
.set_data_type(fe::DataType_t::FLOAT));
auto sdpa_options = fe::graph::SDPA_attributes().set_name("flash_attention");
sdpa_options.set_is_inference(is_inference_only);
sdpa_options.set_attn_scale(attn_scale);
sdpa_options.set_causal_mask(true);
// Create the graph operation and get the output tensors back
auto [O, stats] = graph->sdpa(Q, K, V, sdpa_options);
// Output is (B, T, NH, HS) BF16/FP16 and stats for backward pass is (B, NH, T) FP32
O->set_output(true).set_dim({B, H, T, HS}).set_stride({H * HS * T, HS, H * HS, 1});
assert(stats == nullptr || is_inference_only == false);
if (is_inference_only == false) {
stats->set_output(true).set_data_type(fe::DataType_t::FLOAT)
.set_dim({B, H, T, 1})
.set_stride({H * T, T, 1, 1});
}
assert(graph->validate().is_good());
auto key = graph->key();
auto it = user_maintained_cache_fwd.find(key);
if (it != user_maintained_cache_fwd.end()) {
return it->second;
}
// Build the operation graph and execution part (this is the VERY SLOW PART)
assert(graph->build_operation_graph(cudnn_handle).is_good());
auto plans = graph->create_execution_plans({fe::HeurMode_t::A});
assert(graph->check_support(cudnn_handle).is_good());
assert(graph->build_plans(cudnn_handle).is_good());
auto tuple = std::make_tuple(graph, Q, K, V, attn_scale, O, stats);
user_maintained_cache_fwd.insert({key, tuple});
return tuple;
}
// Used on first run only so we can validate against the CPU results
__global__ void fp32_to_lowp_kernel(floatX* out, const float* inp) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
out[idx] = (floatX)inp[idx];
}
__global__ void lowp_to_fp32_kernel(const floatX* inp, float *out) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
out[idx] = (float)inp[idx];
}
void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS)
float* stats, // output for backward pass: (B, NH, T)
floatX* inp, // input: (B, T, 3, NH, HS) QKV
float* in_fp32, // fp32 input
float* out_fp32, // fp32 output for validation
int B, int T, int C, int NH) {
static bool first_run_validation = true;
int HS = C / NH; // number of features per head
bool is_inference_only = (stats == nullptr);
// Convert from FP32 to FP16/BF16 on 1st run to get correct results
const int block_size = 64; // smallest full occupancy block size on modern GPUs
if (first_run_validation) {
int total_threads = B * T * C * 3;
assert(total_threads % block_size == 0);
int num_blocks = total_threads / block_size;
fp32_to_lowp_kernel<<<num_blocks, block_size>>>(inp, in_fp32);
}
// Get graph and tensors from cache (or generate it on first use)
auto [graph, Q, K, V, attn_scale, O, softmax_stats] =
lookup_cache_or_build_graph_fwd(B, NH, T, HS, is_inference_only);
// Prepare all the tensor pointers for executing the graph
void* devPtrQ = inp;
void* devPtrK = (inp + C);
void* devPtrV = (inp + 2 * C);
float attn_scale_cpu = 1.0 / sqrtf(HS);
void* devPtrO = out;
// Build variant pack
std::unordered_map<std::shared_ptr<fe::graph::Tensor_attributes>, void*> variant_pack = {
{Q, devPtrQ}, {K, devPtrK}, {V, devPtrV}, {attn_scale, &attn_scale_cpu}, {O, devPtrO}};
// Add the stats tensor unless we are only doing inference (only needed for backward pass)
if (is_inference_only == false) {
variant_pack[softmax_stats] = stats;
}
// Reallocate the workspace if the required size is greater than the current workspace
// By default, cuDNN uses up to 256MiB of workspace, so we don't want to just allocate the maximum
if (graph->get_workspace_size() > cudnn_workspace_size) {
if (cudnn_workspace_size > 0) {
cudaCheck(cudaFree(cudnn_workspace));
}
cudnn_workspace_size = graph->get_workspace_size();
cudaCheck(cudaMalloc(&cudnn_workspace, cudnn_workspace_size));
}
// Execute graph
assert(graph->execute(cudnn_handle, variant_pack, cudnn_workspace).is_good());
cudaCheck(cudaGetLastError());
// Optionally convert back from FP16/BF16 to FP32
if (first_run_validation) {
int total_threads = B * T * C;
assert(total_threads % block_size == 0);
int num_blocks = total_threads / block_size;
lowp_to_fp32_kernel<<<num_blocks, block_size>>>(out, out_fp32);
}
cudaCheck(cudaGetLastError());
first_run_validation = false;
}
#endif // ENABLE_CUDNN
// kernel version dispatch
void attention_forward(int kernel_num,
float* out, float* stats, float* vaccum,
float* qkvr, float* preatt, float* att,
float* inp,
int B, int T, int C, int NH,
const int block_size) {
switch (kernel_num) {
case 1:
attention_forward1(out, preatt, att, inp, B, T, C, NH, block_size);
break;
case 2:
attention_forward2(out, inp, B, T, C, NH, block_size);
break;
case 3:
attention_forward3(out, vaccum, qkvr, preatt, att, inp, B, T, C, NH, block_size);
break;
case 4:
attention_forward4(out, vaccum, qkvr, preatt, att, inp, B, T, C, NH, block_size);
break;
case 5:
attention_forward5(out, (floatX*)vaccum, (floatX*)qkvr,
(floatX*)preatt, (floatX*)att,
inp, B, T, C, NH, block_size, false);
break;
case 6: // skip permutes for perf passes (to analyse perf as if in/out were truly 16-bit)
attention_forward5(out, (floatX*)vaccum, (floatX*)qkvr,
(floatX*)preatt, (floatX*)att,
inp, B, T, C, NH, block_size, true);
break;
#ifdef ENABLE_CUDNN
case 10:
// note: validation only cares about out, which is out_fp32 of the function
// inp is hackily converted to FP16 into qkvr only on the first run
// similarly, vaccum is converted to FP32 into out only on the first run
attention_forward_cudnn((floatX*)vaccum, stats, (floatX*)qkvr, inp, out, B, T, C, NH);
break;
#endif
default:
printf("Invalid kernel number\n");
exit(1);
}
}
// ----------------------------------------------------------------------------
int main(int argc, char **argv) {
setup_main();
int B = 8;
int T = 1024;
int C = 768;
int NH = 12;
int deviceIdx = 0;
cudaCheck(cudaSetDevice(deviceIdx));
cudaDeviceProp deviceProp;
cudaGetDeviceProperties(&deviceProp, deviceIdx);
// setup cuBLAS (and cuDNN if needed)
cublasCreate(&cublas_handle);
int enable_tf32 = deviceProp.major >= 8 ? 1 : 0;
printf("enable_tf32: %d\n", enable_tf32);
cublasMath_t cublas_math_mode = enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH;
cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode));
#ifdef ENABLE_CUDNN
checkCudnnErr(cudnnCreate(&cudnn_handle));
#endif
// create host memory of random numbers
float* out = (float*)malloc(B * T * C * sizeof(float));
float* preatt = (float*)malloc(B * NH * T * T * sizeof(float));
float* att = (float*)malloc(B * NH * T * T * sizeof(float));
//float* inp = make_random_float(B * T * 3 * C, 10.0f);
float* inp = make_random_float(B * T * 3 * C);
// move to GPU
float* d_out;
float* d_stats; // for cuDNN
float* d_vaccum;
float* d_qkvr;
float* d_preatt;
float* d_att;
float* d_inp;
cudaCheck(cudaMalloc(&d_out, B * T * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_stats, B * NH * T * sizeof(float)));
cudaCheck(cudaMalloc(&d_vaccum, B * T * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_qkvr, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMalloc(&d_preatt, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMalloc(&d_att, B * NH * T * T * sizeof(float)));
cudaCheck(cudaMalloc(&d_inp, B * T * 3 * C * sizeof(float)));
cudaCheck(cudaMemcpy(d_inp, inp, B * T * 3 * C * sizeof(float), cudaMemcpyHostToDevice));
// read kernel_num from command line
int kernel_num = 1;
if (argc > 1) {
kernel_num = atoi(argv[1]);
}
printf("Using kernel %d\n", kernel_num);
int block_sizes[] = {32, 64, 128, 256, 512};
// Lower accuracy requirements for FP16 (1e-4f also too much for TF32 on kernels 3 & 4)
float accuracy_threshold = (kernel_num <= 4) ? 1e-3f : 1e-2f;
// first check the correctness of the kernel
attention_forward_cpu(out, preatt, att, inp, B, T, C, NH);
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
int block_size = block_sizes[j];
printf("Checking block size %d.\n", block_size);
attention_forward(kernel_num, d_out, d_stats, d_vaccum, d_qkvr, d_preatt, d_att, d_inp, B, T, C, NH, block_size);
// all kernels should produce the correct output out
// todo - make accuracy threshold dynamic and depend on FP16 vs FP32?
validate_result(d_out, out, "out", B * T * C, accuracy_threshold);
// but as for preatt and att, things get a bit more complicated:
if (kernel_num != 2 && kernel_num < 5) {
// kernel 2 (knowingly) fails att/preatt because it uses a different algorithm
// that estimates the softmax online and never materializes preatt/att
validate_result(d_att, att, "att", B * NH * T * T, accuracy_threshold);
}
if (kernel_num != 2 && kernel_num < 4) {
// kernel 4 (knowingly) fails preatt because it fuses the scale normalization
// into the softmax, so preatt is off by 1.0f / sqrt(HS)
// but att and out (checked below) should match.
validate_result(d_preatt, preatt, "preatt", B * NH * T * T, accuracy_threshold);
}
}
printf("All results match. Starting benchmarks.\n\n");
first_run_validation = false;
// benchmark speed of the kernel
for (int j = 0; j < sizeof(block_sizes) / sizeof(int); j++) {
int block_size = block_sizes[j];
int repeat_times = 100;
float elapsed_time = benchmark_kernel(repeat_times, attention_forward,
kernel_num, d_out, d_stats, d_vaccum, d_qkvr, d_preatt, d_att,
d_inp, B, T, C, NH, block_size);
printf("block_size %4d | time %f ms\n", block_size, elapsed_time);
}
// free memory
free(out);
free(preatt);
free(att);
free(inp);
cudaCheck(cudaFree(d_out));
cudaCheck(cudaFree(d_vaccum));
cudaCheck(cudaFree(d_qkvr));
cudaCheck(cudaFree(d_preatt));
cudaCheck(cudaFree(d_att));
cudaCheck(cudaFree(d_inp));
cudaCheck(cudaFree(d_stats));
cublasDestroy(cublas_handle);
#ifdef ENABLE_CUDNN
cudnnDestroy(cudnn_handle);
if (cudnn_workspace_size > 0) {
cudaCheck(cudaFree(cudnn_workspace));
}
#endif
return 0;
}
================================================
FILE: dev/cuda/benchmark_on_modal.py
================================================
"""
Script for running benchmarks on the Modal platform.
This is useful for folks who do not have access to expensive GPUs locally.
Example usage for cuda kernels:
GPU_MEM=80 modal run benchmark_on_modal.py \
--compile-command "nvcc -O3 --use_fast_math attention_forward.cu -o attention_forward -lcublas" \
--run-command "./attention_forward 1"
OR if you want to use cuDNN etc.
For training the gpt2 model with cuDNN use:
GPU_MEM=80 modal run dev/cuda/benchmark_on_modal.py \
--compile-command "make train_gpt2cu USE_CUDNN=1"
--run-command "./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin -j dev/data/tinyshakespeare/tiny_shakespeare_val.bin -v 250 -s 250 -g 144 -f shakespeare.log -b 4"
For profiling using nsight system:
GPU_MEM=80 modal run dev/cuda/benchmark_on_modal.py \
--compile-command "make train_gpt2cu USE_CUDNN=1" \
--run-command "nsys profile --cuda-graph-trace=graph --python-backtrace=cuda --cuda-memory-usage=true \
./train_gpt2cu -i dev/data/tinyshakespeare/tiny_shakespeare_train.bin \
-j dev/data/tinyshakespeare/tiny_shakespeare_val.bin -v 250 -s 250 -g 144 -f shakespeare.log -b 4"
For more nsys profiling specifics and command options, take a look at: https://docs.nvidia.com/nsight-systems/2024.2/UserGuide/
-> To profile the report using a GUI, download NVIDIA NSight System GUI version (this software can run on all OS, so you download it locally)
NOTE: Currently there is a bug in the profiling using nsight system which produces a unrecognized GPU UUId error on the command line but it
does not actually interfere with the model training and validation. The report (that you download) is still generated and can be viewed from Nsight Systems
"""
import subprocess
import os
import sys
import datetime
import modal
from modal import Image, Stub
GPU_NAME_TO_MODAL_CLASS_MAP = {
"H100": modal.gpu.H100,
"A100": modal.gpu.A100,
"A10G": modal.gpu.A10G,
}
N_GPUS = int(os.environ.get("N_GPUS", 1))
GPU_MEM = int(os.environ.get("GPU_MEM", 40))
GPU_NAME = os.environ.get("GPU_NAME", "A100")
GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, size=str(GPU_MEM) + 'GB')
APP_NAME = "llm.c benchmark run"
image = (
Image.from_registry("totallyvyom/cuda-env:latest-2")
.pip_install("huggingface_hub==0.20.3", "hf-transfer==0.1.5")
.env(
dict(
HUGGINGFACE_HUB_CACHE="/pretrained",
HF_HUB_ENABLE_HF_TRANSFER="1",
TQDM_DISABLE="true",
)
)
.run_commands(
"wget -q https://github.com/Kitware/CMake/releases/download/v3.28.1/cmake-3.28.1-Linux-x86_64.sh",
"bash cmake-3.28.1-Linux-x86_64.sh --skip-license --prefix=/usr/local",
"rm cmake-3.28.1-Linux-x86_64.sh",
"ln -s /usr/local/bin/cmake /usr/bin/cmake",)
.run_commands(
"apt-get install -y --allow-change-held-packages libcudnn8 libcudnn8-dev",
"apt-get install -y openmpi-bin openmpi-doc libopenmpi-dev kmod sudo",
"git clone https://github.com/NVIDIA/cudnn-frontend.git /root/cudnn-frontend",
"cd /root/cudnn-frontend && mkdir build && cd build && cmake .. && make"
)
.run_commands(
"wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin && \
mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600 && \
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub && \
add-apt-repository \"deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/ /\" && \
apt-get update"
).run_commands(
"apt-get install -y nsight-systems-2023.3.3"
)
)
stub = modal.App(APP_NAME)
def execute_command(command: str):
command_args = command.split(" ")
print(f"{command_args = }")
subprocess.run(command_args, stdout=sys.stdout, stderr=subprocess.STDOUT)
@stub.function(
gpu=GPU_CONFIG,
image=image,
allow_concurrent_inputs=4,
container_idle_timeout=900,
mounts=[modal.Mount.from_local_dir("./", remote_path="/root/")],
# Instead of 'cuda-env' put your volume name that you create from 'modal volume create {volume-name}'
# This enables the profiling reports to be saved on the volume that you can download by using:
# 'modal volume get {volume-name} {/output_file_name}
# For example right now, when profiling using this command "nsys profile --trace=cuda,nvtx --cuda-graph-trace=graph --python-backtrace=cuda --cuda-memory-usage=true" you would get your report
# using in a directory in your volume, where the name contains the timestamp unique id.
# This script will generate a "report1_{timestamp} folder in volume"
# and you can download it with 'modal volume get {volume-name} report1_{timestamp}
volumes={"/cuda-env": modal.Volume.from_name("cuda-env")},
)
def run_benchmark(compile_command: str, run_command: str):
execute_command("pwd")
execute_command("ls")
execute_command(compile_command)
execute_command(run_command)
# Use this section if you want to profile using nsight system and install the reports on your volume to be locally downloaded
timestamp = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
execute_command("mkdir report1_" + timestamp)
execute_command("mv /root/report1.nsys-rep /root/report1_" + timestamp + "/")
execute_command("mv /root/report1.qdstrm /root/report1_" + timestamp + "/")
execute_command("mv /root/report1_" + timestamp + "/" + " /cuda-env/")
return None
@stub.local_entrypoint()
def inference_main(compile_command: str, run_command: str):
results = run_benchmark.remote(compile_command, run_command)
return results
================================================
FILE: dev/cuda/classifier_fused.cu
================================================
/* Kernels for fused forward/backward classifier part
This fuses softmax, crossentropy, and logit gradients into a single pass, so we don't have to write unnecessary
(B, T, V) tensors. Such an operation is only possible if `dloss` can be known beforehand, which doesn't seem like
much of a restriction: In pretraining, it is just a constant 1/batch_size tensor, for fine-tuning we might zero
out the input prompt, but that is known in advance.
Compile example:
nvcc -O3 --use_fast_math -lcublas -lcublasLt classifier_fused.cu -o classifier_fused
./classifier_fused 1
./classifier_fused 2
./classifier_fused 3
./classifier_fused 4
*/
#include <stdio.h>
#include <stdlib.h>
#include <float.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include "common.h"
// todo - this file does not properly support anything but FP32
// kernel 5 can be run in fp16/bf16 to test performance, but the outputs will be wrong
#if defined(ENABLE_BF16)
typedef __nv_bfloat16 floatX;
#elif defined(ENABLE_FP16)
typedef half floatX;
#else
typedef float floatX;
#endif
typedef Packed128<floatX> x128;
// ----------------------------------------------------------------------------
// CPU code reference
void softmax_forward_cpu(float* out, const float* inp, int N, int C) {
// inp is (N, C)
// out is (N, C), each row of inp will get softmaxed
for (int64_t i = 0; i < N; i++) {
const float* inp_row = inp + i * C;
float* out_row = out + i * C;
float maxval = -INFINITY;
for (int j = 0; j < C; j++) {
if (inp_row[j] > maxval) {
maxval = inp_row[j];
}
}
double sum = 0.0;
for (int j = 0; j < C; j++) {
out_row[j] = expf(inp_row[j] - maxval);
sum += out_row[j];
}
for (int j = 0; j < C; j++) {
out_row[j] /= sum;
}
}
}
void crossentropy_forward_cpu(float* losses,
const float* probs, const int* targets,
int B, int T, int V) {
// output: losses is (B,T) of the individual losses at each position
// input: probs are (B,T,V) of the probabilities
// input: targets is (B,T) of integers giving the correct index in logits
for (int64_t bt = 0; bt < B * T; bt++) {
// loss = -log(probs[target])
const float* probs_bt = probs + bt * V;
int ix = targets[bt];
losses[bt] = -logf(probs_bt[ix]);
}
}
void crossentropy_softmax_backward_cpu(float* dlogits,
const float* dlosses, const float* probs, const int* targets,
int B, int T, int V) {
// backwards through both softmax and crossentropy
for (int64_t bt = 0; bt < B * T; bt++) {
float* dlogits_bt = dlogits + bt * V;
const float* probs_bt = probs + bt * V;
float dloss = dlosses[bt];
int ix = targets[bt];
for (int i = 0; i < V; i++) {
float p = probs_bt[i];
float indicator = i == ix ? 1.0f : 0.0f;
dlogits_bt[i] = (p - indicator) * dloss;
}
}
}
// ----------------------------------------------------
// Kernel Utils
// warp-level reduction for finding the maximum value
__device__ float warpReduceMax(float val) {
for (int offset = 16; offset > 0; offset /= 2) {
val = fmaxf(val, __shfl_xor_sync(0xFFFFFFFF, val, offset));
}
return val;
}
// ----------------------------------------------------------------------------
// GPU kernels
struct SoftmaxParams {
float Scale;
float Offset;
};
namespace cg = cooperative_groups;
__device__ SoftmaxParams prepare_softmax(cg::thread_block_tile<32>& warp,
int64_t idx, const float* inp, int V, int P) {
// this warp (of 32) threads processes one row of inp, i.e. inp[idx, :] of shape (V,)
// note that inp is actually (B * T, P) but we only use the first V elements
// this function then calculates:
// 1) the max value to subtract for numerical stability and
// 2) the sum normalization factor
const float* x = inp + idx * P;
// thread coarsening loop, where the 32 threads serially process all V elements
// thread_rank() is in [0, 31], warp.size() is 32
float maxval = -INFINITY;
float sumval = 0.0f;
for (int i = warp.thread_rank(); i < V; i += warp.size()) {
float v = x[i];
float old_maxval = maxval;
// online softmax recurrence from "Online normalizer calculation for softmax" paper
maxval = fmaxf(maxval, v);
sumval *= expf((old_maxval - maxval));
sumval += expf(v - maxval);
}
// warp-level reduction to get the maxval across the 32 threads
float global_maxval = cg::reduce(warp, maxval, cg::greater<float>{});
// all 32 threads do a final shift of the sum considering the global max in this row
sumval *= expf((maxval - global_maxval));
// warp-level reduction to get the sumval across the 32 threads
float global_sumval = cg::reduce(warp, sumval, cg::plus<float>{});
// the final normalization factor
float norm = 1.0f / global_sumval;
return SoftmaxParams{norm, global_maxval};
}
__global__ void fused_classifier_kernel1(float* dlogits, float* losses,
const float* logits, 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);
// example: B = 4, T = 1024, block_size = 128 => we'd have grid_size = 1024
// each block of 4 warps is in charge of 4 rows of the input, one warp per row
// meta_group_size is the number of warps per block (e.g. 4)
// meta_group_rank is the index of the warp in the block (e.g. 0, 1, 2, 3)
int64_t idx = blockIdx.x * warp.meta_group_size() + warp.meta_group_rank();
if (idx >= B * T) { // there are B * T rows in the input
return;
}
int b = idx / T;
int t = idx % T;
// calculate the offset (maxval) and scale (sumval) for the softmax
SoftmaxParams sp = prepare_softmax(warp, idx, logits, V, P);
// in each row (handled by one warp), thread 0 calculates the loss
// calculate the probability needed for the loss and update losses
if(warp.thread_rank() == 0) {
int ix = targets[b * T + t];
float prob = expf(logits[idx * P + ix] - sp.Offset) * sp.Scale;
losses[b * T + t] = -logf(prob);
}
// finally all threads calculate the gradients
// prob is only materialized here temporarily and in registers, never
// as a full tensor that gets written to global memory
for (int i = warp.thread_rank(); i < V; i += warp.size()) {
float prob = expf(logits[idx * P + i] - sp.Offset) * sp.Scale;
float* dlogits_bt = dlogits + b * T * P + t * P;
float dloss = dlosses[b * T + t];
int ix = targets[b * T + t];
float indicator = i == ix ? 1.0f : 0.0f;
dlogits_bt[i] = (prob - indicator) * dloss;
}
}
__device__ float vec_at(const float4& vec, int index) {
return reinterpret_cast<const float*>(&vec)[index];
}
__device__ SoftmaxParams prepare_softmax_blockwide(cg::thread_block_tile<32>& warp,
int64_t idx, const float* inp, int V, int P) {
// one row of inp, i.e. inp[idx, :] of shape (V,)
// float4 to get 128-bit loads and memory level parallelism
const float4* x_vec4 = reinterpret_cast<const float4*>(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 = ceil_div(V, 4) + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {
float4 v4 = x_vec4[i];
#pragma unroll
for(int k = 0; k < 4; k++) {
if (i*4+k >= V) { // bounds checking against real V
continue;
}
float old_maxval = thread_maxval;
thread_maxval = fmaxf(thread_maxval, vec_at(v4, k));
thread_sumval *= expf(old_maxval - thread_maxval);
thread_sumval += expf(vec_at(v4, k) - 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<float>{});
// 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<float>{});
// 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<float>{});
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<float>{});
// return the softmax parameters
return SoftmaxParams{1.f / block_sumval, block_maxval};
}
// Fused forward and backward pass for classifier including softmax, and logit gradients
// Writes to both probs (only for debugging) and dlogits (only for training) are optional
// N.B.: We may want to reuse the logits memory for dlogits, so they should *not* be __restrict__!
__global__ void fused_classifier_kernel2(float* dlogits, float* losses, float* probs,
const float* logits, 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);
int64_t 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(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 float4* logits_vec4 = reinterpret_cast<const float4*>(logits + idx * P);
for (int i = threadIdx.x; i < ceil_div(V, 4); 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
float4 v4 = __ldcs(&logits_vec4[i]);
#pragma unroll
for(int k = 0; k < 4; ++k) {
int element = i*4 + k;
float prob = expf(vec_at(v4, k) - sp.Offset) * sp.Scale;
prob = (element < V) ? prob : 0.0f; // bounds checking against real V
// this kernel is DRAM limited so cost of inner branch is ~zero
if (probs != NULL) {
probs[idx * P + element] = prob;
}
if (dlogits != NULL) {
float indicator = element == ix ? 1.0f : 0.0f;
dlogits[idx * P + element] = (prob - indicator) * dloss;
}
}
}
}
__device__ SoftmaxParams prepare_softmax_blockwide_nofloat4(cg::thread_block_tile<32>& warp,
int64_t 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<float>{});
// 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<float>{});
// 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<float>{});
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<float>{});
// return the softmax parameters
return SoftmaxParams{1.f / block_sumval, block_maxval};
}
// same as 2 but not using float4
__global__ void fused_classifier_kernel3(float* dlogits, float* losses, float* probs,
const float* logits, 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);
int64_t 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;
}
if (dlogits != NULL) {
float indicator = (i == ix) ? 1.0f : 0.0f;
dlogits[idx * P + i] = (prob - indicator) * dloss;
}
}
}
__device__ SoftmaxParams prepare_softmax_blockwide2(int64_t idx, const floatX* inp, int V, int P) {
// one row of inp, i.e. inp[idx, :] of shape (V,)
const floatX* 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 = ceil_div(V, x128::size) + threadIdx.x - blockDim.x; i >= 0; i -= blockDim.x) {
x128 packed_x = load128cs(x + i * x128::size); // load and do not keep in cache
for(int k = 0; k < packed_x.size; ++k) {
if (i*x128::size+k >= V) { // bounds checking against real V
continue;
}
float v = (float)packed_x[k];
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 = warpReduceMax(thread_maxval);
// 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 = warpReduceMax(warp_maxval);
// 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 = warpReduceSum(thread_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});
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 = warpReduceSum(warp_sumval); //cg::reduce(warp, thread_sumval, cg::plus<float>{});
// return the softmax parameters
return SoftmaxParams{1.f / block_sumval, block_maxval};
}
// same as 2 but using x128
__global__ void fused_classifier_kernel4(floatX* dlogits, floatX* losses, floatX* probs,
const floatX* logits, const floatX* dlosses, const int* targets,
int B, int T, int V, int P) {
int64_t 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_blockwide2(idx, logits, V, P);
// calculate the probability needed for the loss and update (single-threaded)
if(threadIdx.x == 0) {
float prob = expf((float)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 ? (float)dlosses[idx] : 1.0f / (B*T);
// calculate the gradients
gitextract_pp0afoji/ ├── .github/ │ └── workflows/ │ ├── ci.yml │ ├── ci_gpu.yml │ └── ci_tests.yml ├── .gitignore ├── LICENSE ├── Makefile ├── README.md ├── dev/ │ ├── cpu/ │ │ └── matmul_forward.c │ ├── cuda/ │ │ ├── Makefile │ │ ├── README.md │ │ ├── adamw.cu │ │ ├── attention_backward.cu │ │ ├── attention_forward.cu │ │ ├── benchmark_on_modal.py │ │ ├── classifier_fused.cu │ │ ├── common.h │ │ ├── crossentropy_forward.cu │ │ ├── crossentropy_softmax_backward.cu │ │ ├── encoder_backward.cu │ │ ├── encoder_forward.cu │ │ ├── fused_residual_forward.cu │ │ ├── gelu_backward.cu │ │ ├── gelu_forward.cu │ │ ├── global_norm.cu │ │ ├── layernorm_backward.cu │ │ ├── layernorm_forward.cu │ │ ├── matmul_backward.cu │ │ ├── matmul_backward_bias.cu │ │ ├── matmul_forward.cu │ │ ├── nccl_all_reduce.cu │ │ ├── permute.cu │ │ ├── residual_forward.cu │ │ ├── softmax_forward.cu │ │ └── trimat_forward.cu │ ├── data/ │ │ ├── README.md │ │ ├── data_common.py │ │ ├── edu_fineweb.sh │ │ ├── fineweb.py │ │ ├── fineweb.sh │ │ ├── hellaswag.py │ │ ├── mmlu.py │ │ ├── tinyshakespeare.py │ │ └── tinystories.py │ ├── download_starter_pack.sh │ ├── eval/ │ │ ├── README.md │ │ ├── export_hf.py │ │ ├── run_eval.sh │ │ └── summarize_eval.py │ ├── loss_checker_ci.py │ ├── test/ │ │ ├── Makefile │ │ ├── device_file_io.cu │ │ ├── test_dataloader.c │ │ └── test_outlier_detector.c │ ├── unistd.h │ └── vislog.ipynb ├── doc/ │ └── layernorm/ │ ├── layernorm.c │ ├── layernorm.md │ └── layernorm.py ├── llmc/ │ ├── adamw.cuh │ ├── attention.cuh │ ├── cublas_common.h │ ├── cuda_common.h │ ├── cuda_utils.cuh │ ├── cudnn_att.cpp │ ├── cudnn_att.h │ ├── dataloader.h │ ├── encoder.cuh │ ├── fused_classifier.cuh │ ├── gelu.cuh │ ├── global_norm.cuh │ ├── layernorm.cuh │ ├── logger.h │ ├── matmul.cuh │ ├── mfu.h │ ├── outlier_detector.h │ ├── rand.h │ ├── sampler.h │ ├── schedulers.h │ ├── tokenizer.h │ ├── utils.h │ └── zero.cuh ├── profile_gpt2.cu ├── profile_gpt2cu.py ├── requirements.txt ├── scripts/ │ ├── README.md │ ├── multi_node/ │ │ ├── run_gpt2_124M_fs.sbatch │ │ ├── run_gpt2_124M_mpi.sh │ │ └── run_gpt2_124M_tcp.sbatch │ ├── pyrun_gpt2_124M.sh │ ├── run_gpt2_124M.sh │ ├── run_gpt2_1558M.sh │ ├── run_gpt2_350M.sh │ ├── run_gpt2_774M.sh │ └── run_gpt3_125M.sh ├── test_gpt2.c ├── test_gpt2.cu ├── test_gpt2_fp32.cu ├── train_gpt2.c ├── train_gpt2.cu ├── train_gpt2.py ├── train_gpt2_fp32.cu └── train_llama3.py
SYMBOL INDEX (287 symbols across 32 files)
FILE: dev/cpu/matmul_forward.c
function matmul_forward_cpu (line 21) | void matmul_forward_cpu(float* out,
function matmul_forward_ngc92 (line 43) | void matmul_forward_ngc92(float* out,
function matmul_forward (line 92) | void matmul_forward(int kernel_num,
function main (line 114) | int main(int argc, char **argv) {
function validate_results_cpu (line 196) | void validate_results_cpu(const float* kernel_result, const float* cpu_r...
FILE: dev/cuda/benchmark_on_modal.py
function execute_command (line 83) | def execute_command(command: str):
function run_benchmark (line 103) | def run_benchmark(compile_command: str, run_command: str):
function inference_main (line 119) | def inference_main(compile_command: str, run_command: str):
FILE: dev/cuda/common.h
function __device__ (line 16) | __device__ float warpReduceSum(float val) {
function blockReduce (line 30) | float blockReduce(float val, bool final_sync, float out_of_bounds) {
function blockReduce (line 52) | float blockReduce(float val) {
function cuda_check (line 60) | void cuda_check(cudaError_t error, const char *file, int line) {
function cublasCheck (line 70) | void cublasCheck(cublasStatus_t status, const char *file, int line)
function __device__ (line 113) | __device__ explicit Packed128(int4 bits) {
function __device__ (line 118) | __device__ static Packed128 constant(ElementType value) {
function __device__ (line 126) | __device__ static Packed128 zeros() {
function __device__ (line 130) | __device__ static Packed128 ones() {
function __device__ (line 134) | __device__ ElementType& operator[](int index) {
function __device__ (line 137) | __device__ const ElementType& operator[](int index) const {
function __device__ (line 140) | __device__ int4 get_bits() const {
type __nv_bfloat16 (line 186) | typedef __nv_bfloat16 floatX;
type __nv_bfloat16 (line 187) | typedef __nv_bfloat16 floatN;
type half (line 194) | typedef half floatX;
type half (line 195) | typedef half floatN;
type floatX (line 199) | typedef float floatX;
type floatN (line 200) | typedef float floatN;
type Packed128 (line 203) | typedef Packed128<floatX> x128;
function __device__ (line 211) | __device__ floatX __ldcs(const floatX* address) {
FILE: dev/data/data_common.py
function download_file (line 10) | def download_file(url: str, fname: str, chunk_size=1024):
function write_datafile (line 39) | def write_datafile(filename, toks, model_desc="gpt-2"):
function write_evalfile (line 62) | def write_evalfile(filename, datas):
FILE: dev/data/fineweb.py
function tokenize_llama (line 67) | def tokenize_llama(doc):
function tokenize_gpt2 (line 79) | def tokenize_gpt2(doc):
FILE: dev/data/hellaswag.py
function download (line 52) | def download(split):
function render_example (line 63) | def render_example(example):
function iterate_examples (line 102) | def iterate_examples(split):
function evaluate (line 111) | def evaluate(model_type, device):
FILE: dev/data/mmlu.py
function download (line 30) | def download():
function iterate_examples (line 42) | def iterate_examples():
function render_example (line 61) | def render_example(example):
function evaluate (line 90) | def evaluate(model_type, device):
FILE: dev/data/tinyshakespeare.py
function download (line 35) | def download():
function tokenize (line 47) | def tokenize(model_desc):
FILE: dev/data/tinystories.py
function download (line 43) | def download():
function process_shard (line 73) | def process_shard(shard_index, shard_filename, model_desc):
function tokenize (line 98) | def tokenize(model_desc):
FILE: dev/eval/export_hf.py
function tensor_bf16 (line 24) | def tensor_bf16(data_int16, transpose=False):
function tensor_fp32 (line 29) | def tensor_fp32(data_float32, transpose=False):
function convert (line 37) | def convert(filepath, output, push_to_hub=False, out_dtype="bfloat16"):
function spin (line 146) | def spin(output):
FILE: dev/loss_checker_ci.py
function read_numbers_from_file (line 7) | def read_numbers_from_file(file_path, col_start, col_end):
function compare_numbers (line 32) | def compare_numbers(read_values, fixed_values, percent_accuracy):
function main (line 44) | def main():
FILE: dev/test/test_dataloader.c
function check_range (line 18) | void check_range(const int *tokens, const int start, const int end, cons...
function check_equals (line 35) | void check_equals(const int *tokens, const int n, const int expected, co...
function test_simple (line 51) | void test_simple(void) {
function test_multiprocess_simple (line 87) | void test_multiprocess_simple(void) {
function test_shuffled (line 127) | void test_shuffled(void) {
function test_multiprocess_shuffled (line 194) | void test_multiprocess_shuffled(void) {
function main (line 269) | int main(void) {
FILE: dev/test/test_outlier_detector.c
function main (line 11) | int main(void) {
FILE: dev/unistd.h
function clock_gettime (line 20) | static inline int clock_gettime(int ignore_variable, struct timespec* tv)
type glob_t (line 35) | typedef struct glob_t {
function replace_forward_slashes (line 40) | static inline void replace_forward_slashes(char* str) {
function globfree (line 49) | static inline void globfree(glob_t *pglob) {
function glob (line 56) | static inline int glob(const char* pattern, int ignored_flags, int (*ign...
type dirent (line 114) | typedef struct dirent {
type DIR (line 118) | typedef struct DIR {
function DIR (line 124) | static inline DIR *opendir(const char *name) {
type dirent (line 144) | struct dirent
type dirent (line 145) | struct dirent
function closedir (line 160) | static inline int closedir(DIR *directory) {
FILE: doc/layernorm/layernorm.c
function layernorm_forward (line 9) | void layernorm_forward(float* out, float* mean, float* rstd,
function layernorm_backward (line 46) | void layernorm_backward(float* dinp, float* dweight, float* dbias,
function check_tensor (line 90) | int check_tensor(float *a, float *b, int n, char* label) {
function main (line 105) | int main() {
FILE: doc/layernorm/layernorm.py
class LayerNorm (line 5) | class LayerNorm:
method forward (line 8) | def forward(x, w, b):
method backward (line 21) | def backward(dout, cache):
function write (line 56) | def write(tensor, handle):
FILE: llmc/cublas_common.h
function cublasCheck (line 37) | void cublasCheck(cublasStatus_t status, const char *file, int line)
FILE: llmc/cuda_common.h
function cudaCheck_ (line 52) | inline void cudaCheck_(cudaError_t error, const char *file, int line) {
function cudaFreeCheck (line 62) | void cudaFreeCheck(T** ptr, const char *file, int line) {
type PrecisionMode (line 75) | enum PrecisionMode {
type floatX (line 83) | typedef float floatX;
type half (line 87) | typedef half floatX;
type __nv_bfloat16 (line 90) | typedef __nv_bfloat16 floatX;
function __device__ (line 102) | __device__ floatX __ldcs(const floatX* address) {
function device_to_file (line 130) | inline void device_to_file(FILE* dest, void* src, size_t num_bytes, size...
function file_to_device (line 169) | inline void file_to_device(void* dest, FILE* src, size_t num_bytes, size...
FILE: llmc/cudnn_att.cpp
function cuDNNCheck (line 26) | static void cuDNNCheck(cudnnStatus_t error, const char *file, int line) {
function checkCudnnFE (line 34) | static void checkCudnnFE(const fe::error_object& e, const char *file, in...
type UIDs (line 42) | enum UIDs {
function lookup_cache_or_build_graph_fwd (line 60) | auto lookup_cache_or_build_graph_fwd(int B,int H,int T,int HS, int is_in...
function attention_forward_cudnn (line 222) | void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS)
function attention_backward_cudnn (line 256) | void attention_backward_cudnn(floatX* dqkvr, ...
function create_cudnn (line 290) | void create_cudnn() {
function destroy_cudnn (line 294) | void destroy_cudnn() {
FILE: llmc/dataloader.h
type DataLoader (line 29) | typedef struct {
function dataloader_load_shard_ (line 61) | int64_t dataloader_load_shard_(DataLoader *loader, int shard_index) {
function prepare_intra_shard_indices_ (line 99) | void prepare_intra_shard_indices_(DataLoader *loader) {
function dataloader_reset (line 110) | void dataloader_reset(DataLoader *loader) {
function dataloader_advance_ (line 125) | void dataloader_advance_(DataLoader *loader) {
function dataloader_init (line 142) | void dataloader_init(DataLoader *loader,
function dataloader_load_batch (line 203) | void dataloader_load_batch(DataLoader* loader) {
function dataloader_next_batch (line 222) | void dataloader_next_batch(DataLoader *loader) {
function dataloader_resume (line 232) | void dataloader_resume(DataLoader *loader, size_t current_shard_idx, siz...
function dataloader_free (line 239) | void dataloader_free(DataLoader *loader) {
type EvalLoader (line 274) | typedef struct {
function evalloader_reset (line 298) | void evalloader_reset(EvalLoader *loader) {
function evalloader_init (line 340) | void evalloader_init(EvalLoader *loader,
function evalloader_next_example_ (line 380) | void evalloader_next_example_(EvalLoader *loader, int example_batch_inde...
function evalloader_next_batch (line 449) | void evalloader_next_batch(EvalLoader *loader) {
function evalloader_stat_losses (line 468) | int evalloader_stat_losses(EvalLoader *loader, float* losses) {
function evalloader_free (line 511) | void evalloader_free(EvalLoader *loader) {
FILE: llmc/logger.h
type Logger (line 14) | typedef struct {
function logger_init (line 19) | void logger_init(Logger *logger, const char *log_dir, int process_rank, ...
function logger_log_eval (line 34) | void logger_log_eval(Logger *logger, int step, float val) {
function logger_log_val (line 42) | void logger_log_val(Logger *logger, int step, float val_loss) {
function logger_log_train (line 50) | void logger_log_train(Logger *logger, int step, float train_loss, float ...
FILE: llmc/mfu.h
function nvml_check (line 20) | inline void nvml_check(nvmlReturn_t status, const char *file, int line) {
type PerfData (line 30) | typedef struct {
type GPUEntry (line 48) | typedef struct {
function get_flops_promised (line 97) | float get_flops_promised(const char* device, int precision_mode) {
type GPUUtilInfo (line 154) | struct GPUUtilInfo {
function nvmlDevice_t (line 170) | nvmlDevice_t nvml_get_device() {
function GPUUtilInfo (line 196) | GPUUtilInfo get_gpu_utilization_info() {
function GPUUtilInfo (line 239) | GPUUtilInfo get_gpu_utilization_info() {
FILE: llmc/outlier_detector.h
type OutlierDetector (line 18) | typedef struct {
function init_detector (line 26) | void init_detector(OutlierDetector *detector) {
function update_detector (line 36) | double update_detector(OutlierDetector *detector, double new_value) {
FILE: llmc/rand.h
type mt19937_state (line 98) | typedef struct {
function manual_seed (line 106) | void manual_seed(mt19937_state* state, unsigned int seed) {
function next_state (line 118) | void next_state(mt19937_state* state) {
function randint32 (line 134) | unsigned int randint32(mt19937_state* state) {
function randint64 (line 148) | inline unsigned long long randint64(mt19937_state* state) {
function randfloat32 (line 152) | inline float randfloat32(mt19937_state* state) {
function randfloat64 (line 156) | inline double randfloat64(mt19937_state* state) {
function uniform_ (line 160) | void uniform_(float* data, unsigned int numel, float from, float to, mt1...
function normal_fill_16 (line 168) | void normal_fill_16(float* data, float mean, float std) {
function normal_fill (line 180) | void normal_fill(float* data, unsigned int numel, float mean, float std,...
function normal_ (line 197) | void normal_(float* data, unsigned int numel, float mean, float std, mt1...
function init_identity_permutation (line 223) | void init_identity_permutation(int *data, int numel) {
function random_permutation (line 229) | void random_permutation(int* data, int numel, mt19937_state* state) {
FILE: llmc/sampler.h
function random_u32 (line 10) | unsigned int random_u32(unsigned long long *state) {
function random_f32 (line 18) | float random_f32(unsigned long long *state) { // random float32 in [0,1)
function sample_softmax (line 22) | int sample_softmax(const float* logits, int n, float coin) {
FILE: llmc/schedulers.h
type LearningRateScheduler (line 11) | typedef struct {
function lr_scheduler_init (line 19) | void lr_scheduler_init(LearningRateScheduler *scheduler, const char* sch...
function get_learning_rate_cosine (line 28) | float get_learning_rate_cosine(LearningRateScheduler *scheduler, int ste...
function get_learning_rate_linear (line 44) | float get_learning_rate_linear(LearningRateScheduler *scheduler, int ste...
function get_learning_rate_constant (line 58) | float get_learning_rate_constant(LearningRateScheduler *scheduler, int s...
function get_learning_rate_wsd (line 64) | float get_learning_rate_wsd(LearningRateScheduler *scheduler, int step) {
function get_learning_rate (line 83) | float get_learning_rate(LearningRateScheduler *scheduler, int step) {
FILE: llmc/tokenizer.h
type Tokenizer (line 18) | typedef struct {
function safe_printf (line 25) | void safe_printf(const char *piece) {
function tokenizer_init (line 41) | void tokenizer_init(Tokenizer *tokenizer, const char *filename) {
function tokenizer_free (line 98) | void tokenizer_free(Tokenizer *tokenizer) {
FILE: llmc/utils.h
function FILE (line 26) | extern inline FILE *fopen_check(const char *path, const char *mode, cons...
function fread_check (line 44) | extern inline void fread_check(void *ptr, size_t size, size_t nmemb, FIL...
function fclose_check (line 66) | extern inline void fclose_check(FILE *fp, const char *file, int line) {
function sclose_check (line 78) | extern inline void sclose_check(int sockfd, const char *file, int line) {
function closesocket_check (line 91) | extern inline void closesocket_check(int sockfd, const char *file, int l...
function fseek_check (line 104) | extern inline void fseek_check(FILE *fp, long off, int whence, const cha...
function fwrite_check (line 118) | extern inline void fwrite_check(void *ptr, size_t size, size_t nmemb, FI...
function token_check (line 161) | extern inline void token_check(const int* tokens, int token_count, int v...
function create_dir_if_not_exists (line 180) | extern inline void create_dir_if_not_exists(const char *dir) {
function find_max_step (line 192) | extern inline int find_max_step(const char* output_log_dir) {
function ends_with_bin (line 212) | extern inline int ends_with_bin(const char* str) {
FILE: test_gpt2.c
function check_tensor (line 5) | int check_tensor(float *a, float *b, int n, const char* label) {
function main (line 39) | int main(int argc, char *argv[]) {
FILE: train_gpt2.c
function encoder_forward (line 35) | void encoder_forward(float* out,
function encoder_backward (line 60) | void encoder_backward(float* dwte, float* dwpe,
function layernorm_forward (line 78) | void layernorm_forward(float* out, float* mean, float* rstd,
function layernorm_backward (line 120) | void layernorm_backward(float* dinp, float* dweight, float* dbias,
function matmul_forward_naive (line 163) | void matmul_forward_naive(float* out,
function matmul_forward (line 184) | void matmul_forward(float* out,
function matmul_backward (line 231) | void matmul_backward(float* dinp, float* dweight, float* dbias,
function attention_forward (line 271) | void attention_forward(float* out, float* preatt, float* att,
function attention_backward (line 347) | void attention_backward(float* dinp, float* dpreatt, float* datt,
function gelu_forward (line 408) | void gelu_forward(float* out, float* inp, int N) {
function gelu_backward (line 420) | __attribute__((optimize("no-finite-math-only")))
function residual_forward (line 436) | void residual_forward(float* out, float* inp1, float* inp2, int N) {
function residual_backward (line 442) | void residual_backward(float* dinp1, float* dinp2, float* dout, int N) {
function softmax_forward (line 449) | void softmax_forward(float* probs, float* logits, int B, int T, int V, i...
function crossentropy_forward (line 486) | void crossentropy_forward(float* losses,
function crossentropy_softmax_backward (line 502) | void crossentropy_softmax_backward(float* dlogits,
type GPT2Config (line 526) | typedef struct {
type ParameterTensors (line 537) | typedef struct {
function fill_in_parameter_sizes (line 556) | void fill_in_parameter_sizes(size_t* param_sizes, GPT2Config config) {
type ActivationTensors (line 602) | typedef struct {
function fill_in_activation_sizes (line 628) | void fill_in_activation_sizes(size_t* act_sizes, GPT2Config config, int ...
type GPT2 (line 678) | typedef struct {
function gpt2_build_from_checkpoint (line 707) | void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) {
function gpt2_forward (line 765) | void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size...
function gpt2_zero_grad (line 893) | void gpt2_zero_grad(GPT2 *model) {
function gpt2_backward (line 898) | void gpt2_backward(GPT2 *model) {
function gpt2_update (line 1007) | void gpt2_update(GPT2 *model, float learning_rate, float beta1, float be...
function gpt2_free (line 1035) | void gpt2_free(GPT2 *model) {
function random_u32 (line 1051) | unsigned int random_u32(uint64_t *state) {
function random_f32 (line 1058) | float random_f32(uint64_t *state) { // random float32 in [0,1)
function sample_mult (line 1062) | int sample_mult(float* probabilities, int n, float coin) {
function main (line 1077) | int main() {
FILE: train_gpt2.py
class NewGELU (line 40) | class NewGELU(nn.Module):
method forward (line 42) | def forward(self, input):
class CausalSelfAttention (line 48) | class CausalSelfAttention(nn.Module):
method __init__ (line 50) | def __init__(self, config):
method forward (line 65) | def forward(self, x):
class MLP (line 88) | class MLP(nn.Module):
method __init__ (line 90) | def __init__(self, config):
method forward (line 97) | def forward(self, x):
class Block (line 103) | class Block(nn.Module):
method __init__ (line 105) | def __init__(self, config):
method forward (line 112) | def forward(self, x):
class GPTConfig (line 121) | class GPTConfig:
class GPT (line 128) | class GPT(nn.Module):
method __init__ (line 130) | def __init__(self, config):
method _init_weights (line 149) | def _init_weights(self, module):
method forward (line 162) | def forward(self, idx, targets=None, return_logits=True):
method from_pretrained (line 193) | def from_pretrained(cls, model_type):
method configure_optimizers (line 241) | def configure_optimizers(self, weight_decay, learning_rate, betas, dev...
method generate (line 273) | def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
function _peek_data_shard (line 302) | def _peek_data_shard(filename):
function _load_data_shard (line 317) | def _load_data_shard(filename):
class DistributedDataLoader (line 329) | class DistributedDataLoader:
method __init__ (line 330) | def __init__(self, filename_pattern, B, T, process_rank, num_processes):
method reset (line 353) | def reset(self):
method advance (line 361) | def advance(self): # advance to next data shard
method next_batch (line 366) | def next_batch(self):
function write_fp32 (line 383) | def write_fp32(tensor, file):
function write_bf16 (line 388) | def write_bf16(tensor, file):
function write_tensors (line 395) | def write_tensors(model_tensors, L, file, dtype):
function pad_vocab (line 429) | def pad_vocab(tensor, multiple=128, value=0):
function write_model (line 449) | def write_model(model, filename, dtype):
function write_state (line 479) | def write_state(model, x, y, logits, loss, filename):
function write_tokenizer (line 509) | def write_tokenizer(enc, filename):
function print0 (line 529) | def print0(*args, **kwargs):
function get_lr (line 716) | def get_lr(it):
FILE: train_llama3.py
function repeat_kv (line 59) | def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
function reshape_for_broadcast (line 73) | def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
function apply_scaling (line 80) | def apply_scaling(freqs: torch.Tensor):
function apply_rotary_emb (line 104) | def apply_rotary_emb(
function precompute_freqs_cis (line 116) | def precompute_freqs_cis(
class RMSNorm (line 133) | class RMSNorm(torch.nn.Module):
method __init__ (line 134) | def __init__(self, dim: int, eps: float = 1e-6):
method _norm (line 139) | def _norm(self, x):
method forward (line 142) | def forward(self, x):
class CausalSelfAttention (line 146) | class CausalSelfAttention(nn.Module):
method __init__ (line 148) | def __init__(self, config):
method forward (line 167) | def forward(self, x, freqs_cis=None, start_pos=None, mask=None):
class MLP (line 205) | class MLP(nn.Module):
method __init__ (line 207) | def __init__(self, config):
method forward (line 219) | def forward(self, x):
class Block (line 228) | class Block(nn.Module):
method __init__ (line 230) | def __init__(self, config):
method forward (line 237) | def forward(self, x, freqs_cis=None, start_pos=None, mask=None):
class LlamaConfig (line 246) | class LlamaConfig:
method __init__ (line 263) | def __init__(self, **kwargs):
class LLaMA (line 271) | class LLaMA(nn.Module):
method __init__ (line 273) | def __init__(self, config):
method forward (line 295) | def forward(self, idx, targets=None, return_logits=True, start_pos=0):
method adapt_llama_state_dict_keys (line 325) | def adapt_llama_state_dict_keys(checkpoint, config: LlamaConfig):
method adapt_llama_state_dict_keys_hf (line 361) | def adapt_llama_state_dict_keys_hf(checkpoint, config: LlamaConfig):
method from_pretrained_llama3_hf (line 404) | def from_pretrained_llama3_hf(cls, model_id):
method from_pretrained_llama3_meta (line 426) | def from_pretrained_llama3_meta(cls, ckpt_dir, tokenizer_path):
method configure_optimizers (line 444) | def configure_optimizers(self, weight_decay, learning_rate, betas, dev...
method generate (line 476) | def generate(
function sample_top_p (line 559) | def sample_top_p(probs, p):
class Tokenizer (line 596) | class Tokenizer:
method __init__ (line 607) | def __init__(self, model_path: str):
method encode (line 661) | def encode(
method decode (line 717) | def decode(self, t: Sequence[int]) -> str:
method _split_whitespaces_or_nonwhitespaces (line 722) | def _split_whitespaces_or_nonwhitespaces(
function _peek_data_shard (line 750) | def _peek_data_shard(filename):
function _load_data_shard (line 762) | def _load_data_shard(filename):
class DistributedShardedDataLoader (line 774) | class DistributedShardedDataLoader:
method __init__ (line 783) | def __init__(self, filename_pattern, B, T, process_rank, num_processes):
method reset (line 806) | def reset(self):
method advance (line 814) | def advance(self): # advance to next data shard
method next_batch (line 819) | def next_batch(self):
function write_fp32 (line 836) | def write_fp32(tensor, file):
function write_bf16 (line 841) | def write_bf16(tensor, file):
function write_tensors (line 848) | def write_tensors(model_tensors, L, file, dtype):
function write_model (line 870) | def write_model(model, filename, dtype):
function write_state (line 903) | def write_state(model, x, y, logits, loss, filename):
function print0 (line 930) | def print0(*args, **kwargs):
function get_lr (line 1109) | def get_lr(it):
Condensed preview — 102 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (1,286K chars).
[
{
"path": ".github/workflows/ci.yml",
"chars": 9440,
"preview": "name: Build and test\n\non:\n create:\n workflow_dispatch:\n push:\n branches:\n - master\n pull_request:\n branch"
},
{
"path": ".github/workflows/ci_gpu.yml",
"chars": 3638,
"preview": "name: GPU Builds and Tests\n\non:\n create:\n workflow_dispatch:\n push:\n branches:\n - master\n pull_request:\n "
},
{
"path": ".github/workflows/ci_tests.yml",
"chars": 3360,
"preview": "name: Unit, Static and other Tests\n\non:\n create:\n workflow_dispatch:\n push:\n branches:\n - master\n pull_reque"
},
{
"path": ".gitignore",
"chars": 518,
"preview": "# dot files and such\n.vscode\n.venv\n\n# .bin files generated by Python\n*.bin\n\n# data directories\ndev/data/__pycache__/\ndev"
},
{
"path": "LICENSE",
"chars": 1072,
"preview": "MIT License\n\nCopyright (c) 2024 Andrej Karpathy\n\nPermission is hereby granted, free of charge, to any person obtaining a"
},
{
"path": "Makefile",
"chars": 10848,
"preview": "CC ?= clang\nCFLAGS = -Ofast -Wno-unused-result -Wno-ignored-pragmas -Wno-unknown-attributes\nLDFLAGS =\nLDLIBS = -lm\nINCLU"
},
{
"path": "README.md",
"chars": 16466,
"preview": "# llm.c\n\nLLMs in simple, pure C/CUDA with no need for 245MB of PyTorch or 107MB of cPython. Current focus is on pretrain"
},
{
"path": "dev/cpu/matmul_forward.c",
"chars": 7164,
"preview": "/*\nCPU Kernels for matmul forward pass.\n*/\n\n// Compile Examples:\n//\n// MSVC: cl.exe /O2 /fp:fast /Qvec-report:2 /I."
},
{
"path": "dev/cuda/Makefile",
"chars": 3328,
"preview": "# Makefile for building dev/cuda kernels\n# Collects all the make commands in one file but each file also\n# has the compi"
},
{
"path": "dev/cuda/README.md",
"chars": 2347,
"preview": "# dev/cuda\n\nThis directory is scratch space for developing various versions of the needed CUDA kernels. Each file develo"
},
{
"path": "dev/cuda/adamw.cu",
"chars": 9720,
"preview": "/*\nKernels for the AdamW optimizer.\n\nReferences:\n * https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html\n "
},
{
"path": "dev/cuda/attention_backward.cu",
"chars": 49193,
"preview": "/*\nKernels for attention backward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt attention_backwar"
},
{
"path": "dev/cuda/attention_forward.cu",
"chars": 54669,
"preview": "/*\nKernels for attention forward pass.\n\nIf you do not have CUDNN, you can remove ENABLE_CUDNN to run the other kernels\n\n"
},
{
"path": "dev/cuda/benchmark_on_modal.py",
"chars": 5745,
"preview": "\"\"\"\nScript for running benchmarks on the Modal platform.\nThis is useful for folks who do not have access to expensive GP"
},
{
"path": "dev/cuda/classifier_fused.cu",
"chars": 35115,
"preview": "/* Kernels for fused forward/backward classifier part\nThis fuses softmax, crossentropy, and logit gradients into a sing"
},
{
"path": "dev/cuda/common.h",
"chars": 13921,
"preview": "#include <stdlib.h>\n#include <stdio.h>\n#include <cuda_runtime.h>\n#include <cublas_v2.h>\n#include <cublasLt.h>\n#include <"
},
{
"path": "dev/cuda/crossentropy_forward.cu",
"chars": 5045,
"preview": "/*\nKernels for crossentropy forward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt crossentropy_fo"
},
{
"path": "dev/cuda/crossentropy_softmax_backward.cu",
"chars": 6029,
"preview": "/*\nKernels for crossentropy forward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt crossentropy_so"
},
{
"path": "dev/cuda/encoder_backward.cu",
"chars": 6655,
"preview": "/*\nKernels for the positional encoder forward pass in GPT-2.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcubla"
},
{
"path": "dev/cuda/encoder_forward.cu",
"chars": 8567,
"preview": "/*\nKernels for the positional encoder forward pass in GPT-2.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcubla"
},
{
"path": "dev/cuda/fused_residual_forward.cu",
"chars": 28022,
"preview": "/*\nKernels for residual forward pass fused with layernorm\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt"
},
{
"path": "dev/cuda/gelu_backward.cu",
"chars": 6824,
"preview": "/*\nKernels for gelu backward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt gelu_backward.cu -o ge"
},
{
"path": "dev/cuda/gelu_forward.cu",
"chars": 5623,
"preview": "/*\nKernels for gelu forward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt gelu_forward.cu -o gelu"
},
{
"path": "dev/cuda/global_norm.cu",
"chars": 10901,
"preview": "/*\nKernels for a global norm.\nGlobal norm in this context means that we want to calculate a single norm cooperatively us"
},
{
"path": "dev/cuda/layernorm_backward.cu",
"chars": 71946,
"preview": "/*\nKernels for layernorm backward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt layernorm_backwar"
},
{
"path": "dev/cuda/layernorm_forward.cu",
"chars": 24118,
"preview": "/*\nKernels for layernorm forward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt layernorm_forward."
},
{
"path": "dev/cuda/matmul_backward.cu",
"chars": 10934,
"preview": "/*\nKernels for matmul backward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt -Xcompiler -fopenmp "
},
{
"path": "dev/cuda/matmul_backward_bias.cu",
"chars": 27255,
"preview": "/*\nKernels for matmul backward pass bias only.\n\nCompile example:\nnvcc -O3 -lcublas -lcublasLt -std=c++17 matmul_backward"
},
{
"path": "dev/cuda/matmul_forward.cu",
"chars": 18230,
"preview": "/*\nKernels for matmul forward pass.\nIt's advised to use OpenMP here because the CPU implementation is fairly slow otherw"
},
{
"path": "dev/cuda/nccl_all_reduce.cu",
"chars": 7604,
"preview": "/*\n\nA simple test of NCCL capabilities.\nFills a vector with 1s on the first GPU, 2s on the second, etc.\nThen aggregates "
},
{
"path": "dev/cuda/permute.cu",
"chars": 6665,
"preview": "/*\nKernels to demonstrate permute operation.\n\nCompile example:\nnvcc -O3 permute.cu -o permute\n\nThe goal is to permute a "
},
{
"path": "dev/cuda/residual_forward.cu",
"chars": 5418,
"preview": "/*\nKernels for residual forward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt residual_forward.cu"
},
{
"path": "dev/cuda/softmax_forward.cu",
"chars": 25049,
"preview": "/*\nKernels for softmax forward pass.\n\nCompile example:\nnvcc -O3 --use_fast_math -lcublas -lcublasLt softmax_forward.cu -"
},
{
"path": "dev/cuda/trimat_forward.cu",
"chars": 27483,
"preview": "/*\nTriangular matrix multiplication as in autoregressive attention. A short story.\nby @ngc92\n\nCompile:\nnvcc -O3 --use_fa"
},
{
"path": "dev/data/README.md",
"chars": 741,
"preview": "# dev/data organization\n\nThe idea is that each dataset has a .py file here in the root of `dev/data`, and each dataset t"
},
{
"path": "dev/data/data_common.py",
"chars": 4908,
"preview": "\"\"\"\nCommon utilities for the datasets\n\"\"\"\n\nimport requests\nfrom tqdm import tqdm\nimport numpy as np\n\n\ndef download_file("
},
{
"path": "dev/data/edu_fineweb.sh",
"chars": 2008,
"preview": "#!/bin/bash\n\n# Downloads the FineWeb-Edu 100B dataset, but in an already tokenized format in .bin files\n# Example: ./edu"
},
{
"path": "dev/data/fineweb.py",
"chars": 6286,
"preview": "\"\"\"\nFineWeb dataset (for srs pretraining)\nhttps://huggingface.co/datasets/HuggingFaceFW/fineweb\n\nexample doc to highligh"
},
{
"path": "dev/data/fineweb.sh",
"chars": 1992,
"preview": "#!/bin/bash\n\n# Downloads the FineWeb100B dataset, but in an already tokenized format in .bin files\n# Example: ./fineweb."
},
{
"path": "dev/data/hellaswag.py",
"chars": 7646,
"preview": "\"\"\"\nDownloads and evaluates HellaSwag in Python.\nThis then acts as the reference file for llm.c\nAlso writes the data (to"
},
{
"path": "dev/data/mmlu.py",
"chars": 5816,
"preview": "\"\"\"\nDownloads and evaluates MMLU in Python.\nThis then acts as the reference file for llm.c\nhttps://github.com/hendrycks/"
},
{
"path": "dev/data/tinyshakespeare.py",
"chars": 4068,
"preview": "\"\"\"\nDownloads and tokenizes the TinyShakespeare dataset.\n- The download is from Github.\n- The tokenization is GPT-2 toke"
},
{
"path": "dev/data/tinystories.py",
"chars": 4925,
"preview": "\"\"\"\nDownloads and tokenizes the TinyStories dataset.\n- The download is from HuggingFace datasets.\n- The tokenization is "
},
{
"path": "dev/download_starter_pack.sh",
"chars": 2046,
"preview": "#!/bin/bash\n\n# Get the directory of the script\nSCRIPT_DIR=$(dirname \"$(realpath \"$0\")\")\n\n# Base URL\nBASE_URL=\"https://hu"
},
{
"path": "dev/eval/README.md",
"chars": 2557,
"preview": "# eleuther eval readme\n\nThe goal here is to run the Eleuther Eval harness exactly in the same way as that used in the [h"
},
{
"path": "dev/eval/export_hf.py",
"chars": 7304,
"preview": "\"\"\"\nScript to convert GPT2 models from llm.c binary format to Hugging Face\n\nIt can optinally upload to your account on H"
},
{
"path": "dev/eval/run_eval.sh",
"chars": 4899,
"preview": "# https://huggingface.co/spaces/open-llm-leaderboard/open_llm_leaderboard\n# (See About tab -> REPRODUCIBILITY)\n\n# This s"
},
{
"path": "dev/eval/summarize_eval.py",
"chars": 1074,
"preview": "# example run command\n# python dev/eval/summarize_eval.py lm-evaluation-harness/results/result774M\n# this script is opti"
},
{
"path": "dev/loss_checker_ci.py",
"chars": 2999,
"preview": "# Description: A script to compare numbers in a file with fixed values and check for accuracy within a specified percent"
},
{
"path": "dev/test/Makefile",
"chars": 5715,
"preview": "CC ?= gcc\n# example: make test_dataloader TEST_CFLAGS=-fsanitize=address -fno-omit-frame-pointer \nCFLAGS = -Ofast -Wno-u"
},
{
"path": "dev/test/device_file_io.cu",
"chars": 1946,
"preview": "/*\nTests device <-> file IO functions\n\ncompile and run as (from dev/test directory)\nnvcc -o device_file_io device_file_i"
},
{
"path": "dev/test/test_dataloader.c",
"chars": 11929,
"preview": "/*\nTests our DataLoader\n\ncompile and run as (from dev/test directory)\ngcc -O3 -I../../llmc -o test_dataloader test_datal"
},
{
"path": "dev/test/test_outlier_detector.c",
"chars": 1726,
"preview": "/*\nTests our OutlierDetector\n\ncompile and run as (from dev/test directory)\ngcc -O3 -I../../llmc -o test_outlier_detector"
},
{
"path": "dev/unistd.h",
"chars": 4995,
"preview": "// header file that is necessary to compile on Windows\n#ifndef UNISTD_H\n#define UNISTD_H\n\n#define _CRT_SECURE_NO_WARNING"
},
{
"path": "dev/vislog.ipynb",
"chars": 5702,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {},\n \"source\": [\n \"Simple visualizer for log files wri"
},
{
"path": "doc/layernorm/layernorm.c",
"chars": 6204,
"preview": "// must run `python layernorm.py` first to generate the reference data\n// then compile for example as `gcc layernorm.c -"
},
{
"path": "doc/layernorm/layernorm.md",
"chars": 18425,
"preview": "\n# layernorm\n\nQuick tutorial. Let's look at how LayerNorm is handled, as one example layer in the model. We start with t"
},
{
"path": "doc/layernorm/layernorm.py",
"chars": 1996,
"preview": "import torch\n\neps = 1e-5\n\nclass LayerNorm:\n\n @staticmethod\n def forward(x, w, b):\n B, T, C = x.size()\n "
},
{
"path": "llmc/adamw.cuh",
"chars": 5477,
"preview": "/*\nAdamW kernel\n*/\n\n// llmc internal imports\n#include \"cuda_common.h\"\n#include \"cuda_utils.cuh\"\n\n// --------------------"
},
{
"path": "llmc/attention.cuh",
"chars": 11324,
"preview": "/*\nAttention, as a fallback when we do not use the Flash Attention from cuDNN\n*/\n#include <assert.h>\n// llmc internal im"
},
{
"path": "llmc/cublas_common.h",
"chars": 1306,
"preview": "/*\ncuBLAS related utils\n*/\n#ifndef CUBLAS_COMMON_H\n#define CUBLAS_COMMON_H\n\n#include <stddef.h>\n#include <stdlib.h>\n#inc"
},
{
"path": "llmc/cuda_common.h",
"chars": 8388,
"preview": "/*\nCommon utilities for CUDA code.\n*/\n#ifndef CUDA_COMMON_H\n#define CUDA_COMMON_H\n\n#include <stdlib.h>\n#include <stdio.h"
},
{
"path": "llmc/cuda_utils.cuh",
"chars": 11247,
"preview": "// Utilities for use in __device__ code\n\n#ifndef CUDA_UTILS_CUH\n#define CUDA_UTILS_CUH\n\n#include \"cuda_common.h\"\n\n// ---"
},
{
"path": "llmc/cudnn_att.cpp",
"chars": 12945,
"preview": "// all cudnn-related functions are in this file, so that they don't need to be recompiled everytime\n// we change some un"
},
{
"path": "llmc/cudnn_att.h",
"chars": 799,
"preview": "/*\ncuDNN (flash) attention\n*/\n#ifndef CUDNN_ATT_H\n#define CUDNN_ATT_H\n\n#include \"cuda_common.h\"\n\n// forward declarations"
},
{
"path": "llmc/dataloader.h",
"chars": 24450,
"preview": "/*\nImplements:\n- DataLoader for model training. Reads and serves data shards.\n- EvalLoader for multiple-choice evaluatio"
},
{
"path": "llmc/encoder.cuh",
"chars": 11302,
"preview": "/*\nThe GPT-2 Encoder, which combines two encodings: token and position\nIn the forward pass, both encodings are added tog"
},
{
"path": "llmc/fused_classifier.cuh",
"chars": 6800,
"preview": "/*\nFused Classifier:\n- Forwards the Cross Entropy Loss\n- Never materializes the full normalized logits, only at the targ"
},
{
"path": "llmc/gelu.cuh",
"chars": 2662,
"preview": "/*\n(Approximate) GeLU non-linearity layer\n*/\n#include <assert.h>\n// llmc internal imports\n#include \"cuda_common.h\"\n#incl"
},
{
"path": "llmc/global_norm.cuh",
"chars": 3715,
"preview": "/*\nGlobal norm, used in gradient clipping\n*/\n#include <assert.h>\n#include <stddef.h>\n#include <cuda_runtime_api.h>\n// ll"
},
{
"path": "llmc/layernorm.cuh",
"chars": 22720,
"preview": "/*\nLayerNorm CUDA kernel, and also Residual, because sometimes they are fused\n\nNote in llm.c we try to be clever in the "
},
{
"path": "llmc/logger.h",
"chars": 1865,
"preview": "/*\nImplements a simple logger that writes log files in the output directory.\nThe Logger object is stateless and uses app"
},
{
"path": "llmc/matmul.cuh",
"chars": 14117,
"preview": "/*\nMatrix Multiplication, with help from cuBLASLt\n*/\n#include <assert.h>\n#include <type_traits> // std::bool_consta"
},
{
"path": "llmc/mfu.h",
"chars": 10317,
"preview": "#ifndef MFU_H\n#define MFU_H\n\n#include <stdio.h>\n#include <stdlib.h>\n#include <string.h>\n#if __has_include(<nvml.h>)\n#def"
},
{
"path": "llmc/outlier_detector.h",
"chars": 2498,
"preview": "/*\nSimple OutlierDetector that we can use to monitor the loss and grad norm\nInternally, it keeps track of a window of me"
},
{
"path": "llmc/rand.h",
"chars": 7453,
"preview": "/*\nMersenne Twisters implementation, numerically identical to torch.\n\nExample usage:\n\n mt19937_state state;\n manua"
},
{
"path": "llmc/sampler.h",
"chars": 1146,
"preview": "/*\nImplements a simple Sampler, used during model inference to sample tokens.\n*/\n#ifndef SAMPLER_H\n#define SAMPLER_H\n\n#i"
},
{
"path": "llmc/schedulers.h",
"chars": 4340,
"preview": "/*\nImplements various learning rate schedulers.\n*/\n#ifndef SCHEDULERS_H\n#define SCHEDULERS_H\n\n#include <assert.h>\n#inclu"
},
{
"path": "llmc/tokenizer.h",
"chars": 3691,
"preview": "/*\nDefines the GPT-2 Tokenizer.\nOnly supports decoding, i.e.: tokens (integers) -> strings\nThis is all we need for uncon"
},
{
"path": "llmc/utils.h",
"chars": 8662,
"preview": "/*\n This file contains utilities shared between the different training scripts.\n In particular, we define a series of ma"
},
{
"path": "llmc/zero.cuh",
"chars": 23059,
"preview": "/*\nUtilities for ZeRO sharding\n*/\n\n#ifndef LLMC_ZERO_CUH\n#define LLMC_ZERO_CUH\n\n#include <cuda_runtime_api.h>\n#include <"
},
{
"path": "profile_gpt2.cu",
"chars": 2948,
"preview": "/*\nThis code is a convenience tool for profiling the CUDA kernels in the training\nloop of train_gpt2.cu. Compile:\n\nmake "
},
{
"path": "profile_gpt2cu.py",
"chars": 8681,
"preview": "# runs profiling with ncu, generates a `profile.ncu-rep` for viewing with NSight Compute, and prints out\n# basic kernel "
},
{
"path": "requirements.txt",
"chars": 59,
"preview": "tqdm\nnumpy<2\ntorch\ntiktoken\ntransformers\ndatasets\nrequests\n"
},
{
"path": "scripts/README.md",
"chars": 2745,
"preview": "# scripts\n\nThese shell scripts hold the exact commands to llm.c that reproduce the GPT-2 and GPT-3 runs.\n\n### pytorch re"
},
{
"path": "scripts/multi_node/run_gpt2_124M_fs.sbatch",
"chars": 3170,
"preview": "#!/bin/bash\n#SBATCH --job-name=llmc-multinode # job name\n#SBATCH --output=/home/ubun"
},
{
"path": "scripts/multi_node/run_gpt2_124M_mpi.sh",
"chars": 1579,
"preview": "\nmake train_gpt2cu USE_CUDNN=1\n\n# NOTE: change the following to match your system\nbinary_path=\"/home/ubuntu/llm.c/train_"
},
{
"path": "scripts/multi_node/run_gpt2_124M_tcp.sbatch",
"chars": 3195,
"preview": "#!/bin/bash\n#SBATCH --job-name=llmc-multinode # job name\n#SBATCH --output=/home/ubun"
},
{
"path": "scripts/pyrun_gpt2_124M.sh",
"chars": 876,
"preview": "#!/bin/bash\n\n# the same as scripts/run_gpt2_124M.sh but with PyTorch\n\n# if you wish to train on just a single GPU, simpl"
},
{
"path": "scripts/run_gpt2_124M.sh",
"chars": 1289,
"preview": "# GPT-2 (124M) repro on FineWeb\n# 124M parameter model on 10B tokens\n# => 6 * 124e6 * 10e9 = 7.44e18 ~= 7e18 capability "
},
{
"path": "scripts/run_gpt2_1558M.sh",
"chars": 1306,
"preview": "# GPT-2 (1558M) repro on FineWeb-EDU\n# 1558M parameter model on 32B tokens\n# => 6 * 1558e6 * 32e9 = 6.966e20 ~= 3e20 cap"
},
{
"path": "scripts/run_gpt2_350M.sh",
"chars": 1352,
"preview": "# GPT-2 (350M) repro on FineWeb\n# 350M parameter model on ~30B tokens\n# => 6 * 350e6 * 31.5e9 = 6.615e19 ~= 7e19 capabil"
},
{
"path": "scripts/run_gpt2_774M.sh",
"chars": 1372,
"preview": "# GPT-2 (774M) repro on FineWeb\n# 774M parameter model on ~150B tokens\n# => 6 * 774e6 * 150e9 = 6.966e20 ~= 7e20 capabil"
},
{
"path": "scripts/run_gpt3_125M.sh",
"chars": 1323,
"preview": "# GPT-3 (125M) repro, but using FineWeb\n# 125M parameter model on 300B tokens\n# note context length: 1024 -> 2048 for GP"
},
{
"path": "test_gpt2.c",
"chars": 8030,
"preview": "#define TESTING\n#include \"train_gpt2.c\"\n\n// poor man's tensor checker\nint check_tensor(float *a, float *b, int n, const "
},
{
"path": "test_gpt2.cu",
"chars": 17190,
"preview": "#define TESTING\n#include \"train_gpt2.cu\"\n\n// poor man's tensor checker\nint check_tensor(float *a, float *b, int n, const"
},
{
"path": "test_gpt2_fp32.cu",
"chars": 11089,
"preview": "#define TESTING\n#include \"train_gpt2_fp32.cu\"\n\n// poor man's tensor checker\nint check_tensor(float *a, float *b, int n, "
},
{
"path": "train_gpt2.c",
"chars": 50680,
"preview": "/*\nThis file trains the GPT-2 model.\nThis version is the clean, minimal, reference. As such:\n- it runs on CPU.\n- it does"
},
{
"path": "train_gpt2.cu",
"chars": 104647,
"preview": "/*\nGPT-2 Transformer Neural Net training loop. See README.md for usage.\n*/\n#include <unistd.h>\n#include <stdio.h>\n#inclu"
},
{
"path": "train_gpt2.py",
"chars": 41860,
"preview": "\"\"\"\nReference code for GPT-2 training and inference.\nWill save the model weights into files, to be read from C as initia"
},
{
"path": "train_gpt2_fp32.cu",
"chars": 77134,
"preview": "/*\nGPT-2 Transformer Neural Net trained in raw CUDA\nNon-trivial notes to be aware of:\n\nWe are being clever in the backwa"
},
{
"path": "train_llama3.py",
"chars": 57538,
"preview": "\"\"\"\nReference code for LLaMA-3.1 training and inference.\nWill save the model weights into files, to be read from C as in"
}
]
About this extraction
This page contains the full source code of the karpathy/llm.c GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 102 files (1.2 MB), approximately 345.8k tokens, and a symbol index with 287 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.