mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-28 20:35:09 -04:00
round 1 of small fixes for cudnn support
This commit is contained in:
Andrej Karpathy
parent
1147983b8c
commit
04718fbbd1
2 changed files with 41 additions and 23 deletions
49
Makefile
49
Makefile
|
|
@ -14,11 +14,14 @@ CUDA_OUTPUT_FILE = -o $@
|
|||
# NVCC flags
|
||||
# -t=0 is short for --threads, 0 = number of CPUs on the machine
|
||||
NVCC_FLAGS = -O3 -t=0 --use_fast_math
|
||||
NVCC_LDFLAGS = -lcublas -lcublasLt -lcudnn
|
||||
NVCC_INCLUDES = -I../cudnn-frontend/include # TODO: Don't hardcode path
|
||||
NVCC_LDFLAGS = -lcublas -lcublasLt
|
||||
NVCC_INCLUDES =
|
||||
NVCC_LDLIBS =
|
||||
NCLL_INCUDES =
|
||||
|
||||
# autodect a lot of various supports on current platform
|
||||
$(info ---------------------------------------------)
|
||||
|
||||
ifneq ($(OS), Windows_NT)
|
||||
NVCC := $(shell which nvcc 2>/dev/null)
|
||||
|
||||
|
|
@ -59,6 +62,21 @@ else
|
|||
endif
|
||||
endif
|
||||
|
||||
# Check and include cudnn if available
|
||||
# Currently hard-coding a bunch of stuff here for Linux, todo make this better/nicer
|
||||
ifeq ($(SHELL_UNAME), Linux)
|
||||
ifeq ($(shell [ -d ../cudnn-frontend/include ] && echo "exists"), exists)
|
||||
$(info ✓ cuDNN found, will run with flash-attention)
|
||||
NVCC_INCLUDES += -I../cudnn-frontend/include
|
||||
NVCC_LDFLAGS += -lcudnn
|
||||
CFLAGS += -DENABLE_CUDNN
|
||||
else
|
||||
$(info ✗ cuDNN not found, disabling flash-attention)
|
||||
$(info ---> You need cuDNN: https://developer.nvidia.com/cudnn)
|
||||
$(info ---> And the cuDNN front-end: https://github.com/NVIDIA/cudnn-frontend/tree/main)
|
||||
endif
|
||||
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
|
||||
|
|
@ -79,46 +97,47 @@ else
|
|||
LDFLAGS += -L/opt/homebrew/opt/libomp/lib
|
||||
LDLIBS += -lomp
|
||||
INCLUDES += -I/opt/homebrew/opt/libomp/include
|
||||
$(info OpenMP found, compiling with OpenMP support)
|
||||
$(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, compiling with OpenMP support)
|
||||
$(info ✓ OpenMP found)
|
||||
else
|
||||
$(warning OpenMP not found, skipping OpenMP support)
|
||||
$(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, compiling with OpenMP support)
|
||||
$(info ✓ OpenMP found)
|
||||
else
|
||||
$(warning OpenMP not found, skipping OpenMP support)
|
||||
$(info ✗ OpenMP not found)
|
||||
endif
|
||||
endif
|
||||
endif
|
||||
endif
|
||||
|
||||
# Check if OpenMPI and NCCL are available, include them if so, for multi-GPU training
|
||||
ifeq ($(NO_MULTI_GPU), 1)
|
||||
$(info Multi-GPU (OpenMPI + NCCL) is manually disabled)
|
||||
$(info → Multi-GPU (OpenMPI + NCCL) is manually disabled)
|
||||
else
|
||||
ifneq ($(OS), Windows_NT)
|
||||
# Detect if running on macOS or Linux
|
||||
ifeq ($(SHELL_UNAME), Darwin)
|
||||
$(warning Multi-GPU on CUDA on Darwin is not supported, skipping OpenMPI + NCCL support)
|
||||
$(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping OpenMPI + NCCL support)
|
||||
else ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists)
|
||||
$(info OpenMPI found, adding support)
|
||||
$(info ✓ OpenMPI found, OK to train with multiple GPUs)
|
||||
NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include
|
||||
NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/
|
||||
NVCC_LDLIBS += -lmpi -lnccl
|
||||
NVCC_FLAGS += -DMULTI_GPU
|
||||
else
|
||||
$(warning OpenMPI is not found, disabling multi-GPU support)
|
||||
$(warning On Linux you can try install OpenMPI with `sudo apt install openmpi-bin openmpi-doc libopenmpi-dev`)
|
||||
$(info ✗ OpenMPI is not found, disabling multi-GPU support)
|
||||
$(info ---> On Linux you can try install OpenMPI with `sudo apt install openmpi-bin openmpi-doc libopenmpi-dev`)
|
||||
endif
|
||||
endif
|
||||
endif
|
||||
|
|
@ -145,12 +164,14 @@ TARGETS = train_gpt2 test_gpt2
|
|||
|
||||
# Conditional inclusion of CUDA targets
|
||||
ifeq ($(NVCC),)
|
||||
$(info nvcc not found, skipping CUDA builds)
|
||||
$(info ✗ nvcc not found, skipping GPU/CUDA builds)
|
||||
else
|
||||
$(info nvcc found, including CUDA builds)
|
||||
$(info ✓ nvcc found, including GPU/CUDA support)
|
||||
TARGETS += train_gpt2cu test_gpt2cu train_gpt2fp32cu test_gpt2fp32cu
|
||||
endif
|
||||
|
||||
$(info ---------------------------------------------)
|
||||
|
||||
all: $(TARGETS)
|
||||
|
||||
train_gpt2: train_gpt2.c
|
||||
|
|
|
|||
|
|
@ -32,7 +32,6 @@ make train_gpt2cu PRECISION=FP32
|
|||
This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200),
|
||||
-a 1 is "overfit single batch", -x 10 is 10 iterations, and -f 0 disables tf32
|
||||
*/
|
||||
#define ENABLE_CUDNN // can be enabled via nvcc "-DENABLE_CUDNN"
|
||||
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
|
|
@ -2474,11 +2473,11 @@ int main(int argc, char *argv[]) {
|
|||
cuda_arch_major = deviceProp.major;
|
||||
cuda_arch_minor = deviceProp.minor;
|
||||
|
||||
// setup cuBLAS and cuBLASLt
|
||||
// set up cuBLAS and cuBLASLt
|
||||
cublasCheck(cublasCreate(&cublas_handle));
|
||||
cublasCheck(cublasLtCreate(&cublaslt_handle));
|
||||
cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size));
|
||||
|
||||
// setup compute precision settings for cublas
|
||||
// TF32 precision is equivalent to torch.set_float32_matmul_precision('high')
|
||||
int enable_tf32 = cuda_arch_major >= 8 ? 1 : 0;
|
||||
if (override_enable_tf32 == 0) { enable_tf32 = 0; } // force to zero via arg
|
||||
|
|
@ -2487,6 +2486,7 @@ int main(int argc, char *argv[]) {
|
|||
cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode));
|
||||
if(cublas_compute_type); // unused in BF16 mode, avoid warning
|
||||
|
||||
// set up cuDNN
|
||||
#ifdef ENABLE_CUDNN
|
||||
checkCudnnErr(cudnnCreate(&cudnn_handle));
|
||||
#endif
|
||||
|
|
@ -2648,7 +2648,7 @@ int main(int argc, char *argv[]) {
|
|||
// add a total average, for optimizations that are only mild improvements (excluding 1st batch as warmup)
|
||||
printf0("total average iteration time: %f ms\n", total_sum_iteration_time_s / (train_num_batches-1) * 1000);
|
||||
|
||||
// free
|
||||
// free and destroy everything
|
||||
dataloader_free(&train_loader);
|
||||
dataloader_free(&val_loader);
|
||||
tokenizer_free(&tokenizer);
|
||||
|
|
@ -2656,13 +2656,10 @@ int main(int argc, char *argv[]) {
|
|||
free(cpu_logits_raw);
|
||||
free(cpu_logits);
|
||||
free(gen_tokens);
|
||||
|
||||
#ifdef ENABLE_CUDNN
|
||||
if (cudnn_workspace != NULL) {
|
||||
cudaCheck(cudaFree(cudnn_workspace));
|
||||
}
|
||||
#endif
|
||||
if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); }
|
||||
checkCudnnErr(cudnnDestroy(cudnn_handle));
|
||||
#endif
|
||||
cudaCheck(cudaFree(cublaslt_workspace));
|
||||
cublasCheck(cublasDestroy(cublas_handle));
|
||||
cublasCheck(cublasLtDestroy(cublaslt_handle));
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue