From 74d5bcd03a5326d57d4d40ee4ba60fbb7bc2e049 Mon Sep 17 00:00:00 2001 From: harry Date: Tue, 30 Apr 2024 14:51:06 -0700 Subject: [PATCH 01/15] Adding a java port of the project --- README.md | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 0eb8dd3..0328eb8 100644 --- a/README.md +++ b/README.md @@ -290,23 +290,27 @@ Lastly, I will be a lot more sensitive to complexity in the root folder of the p ## notable forks -- Mojo - - [llm.🔥](https://github.com/dorjeduck/llm.mojo) by @[dorjeduck](https://github.com/dorjeduck): a Mojo port of this project - - C# - [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): a C# port of this project -- Rust - - [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project +- 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 + +- Rust + - [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust 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 -- Go - - [llm.go](https://github.com/joshcarp/llm.go) by @[joshcarp](https://github.com/joshcarp): a Go port of this project ## discussions From e780b56974ea85a12e300ff26e4b70105434a0f8 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 3 May 2024 15:54:29 +0000 Subject: [PATCH 02/15] small fixes for profiling --- Makefile | 2 +- profile_gpt2.cu | 8 ++++---- profile_gpt2cu.py | 9 ++++++--- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/Makefile b/Makefile index 0692380..3af0225 100644 --- a/Makefile +++ b/Makefile @@ -207,7 +207,7 @@ test_gpt2fp32cu: test_gpt2_fp32.cu $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) profile_gpt2cu: profile_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) clean: $(REMOVE_FILES) $(TARGETS) diff --git a/profile_gpt2.cu b/profile_gpt2.cu index f412eed..5b043e0 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -49,7 +49,6 @@ int main() { cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode)); // setup the (global) cuBLASLt workspace cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); - #ifdef ENABLE_CUDNN checkCudnnErr(cudnnCreate(&cudnn_handle)); #endif @@ -58,8 +57,8 @@ int main() { GPT2 model; gpt2_build_from_checkpoint(&model, "gpt2_124M_bf16.bin"); - int B = 4; - int T = 1024; + int B = 24; // if program OOMs decrease this number, e.g. all the way down to 4 or etc + int T = 1024; // if even that OOMs move on to this one. keep them nice and powers of 2 printf("batch size: %d\n", B); printf("sequence length: %d\n", T); @@ -70,6 +69,7 @@ int main() { y[i] = i % model.config.vocab_size; } + // override number of layers to 1 because all layers repeat the same kernels, only profile once model.config.num_layers = 1; // do a training step @@ -78,9 +78,9 @@ int main() { gpt2_backward(&model); gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, 1); cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings + // free gpt2_free(&model); - #ifdef ENABLE_CUDNN if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } checkCudnnErr(cudnnDestroy(cudnn_handle)); diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index b3eec86..d493687 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -14,13 +14,13 @@ NCU = shutil.which("ncu") if NCU is None: NCU = "/usr/local/cuda/bin/ncu" -# build the exe -subprocess.check_call(["make", "profile_gpt2cu"]) +# build the executable +subprocess.check_call(["make", "profile_gpt2cu", "NO_MULTI_GPU=1", "USE_CUDNN=1"]) # record metrics # --full and --import-source are entirely superfluous for this script, but you might want to # manually inspect `profile.ncu-rep`, so we keep it here -cmd = [NCU, "--set", "full", "--import-source", "yes", "-o", "profile", "-f", "./profile_gpt2cu"] +cmd = ["sudo", NCU, "--set", "full", "--import-source", "yes", "-o", "profile", "-f", "./profile_gpt2cu"] subprocess.check_call(cmd) # generate csv @@ -129,6 +129,9 @@ print("Kernel type summaries:") print(f" {'name':<40} {'time':>6} {'frac':>6}") ordered = sorted(summaries.items(), key=lambda x: x[1], reverse=True) for entry, value in ordered: + # crop entry to be at most 40 characters + if len(entry) > 40: + entry = entry[:37] + "..." print(f" {entry:<40} {value:6.2f} {100*value / total_time:6.2f}%") From 325b45647f2458b5aa6af2981d4d7beb1e88b17f Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 3 May 2024 21:14:39 +0300 Subject: [PATCH 03/15] updated the profile script to be more robust, and adapted to recent kernel changes make it compile with cudnn default batchsize is 24 now --- profile_gpt2cu.py | 82 ++++++++++++++++++++++++++++++++++------------- 1 file changed, 59 insertions(+), 23 deletions(-) diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index d493687..448dd5a 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -39,32 +39,45 @@ result = subprocess.check_output(cmd, text=True).strip() reader = csv.reader(result.splitlines(keepends=True)) # model config -CLS_START = 15 +CLS_START = -1 CLS_NUM = 6 -ADAM_ID = 44 N_LAYERS = 12 summaries = defaultdict(lambda: 0.0) +counts = defaultdict(lambda: 0) passes = defaultdict(lambda: 0.0) total = defaultdict(lambda: 0.0) no_cutlass = 0.0 CC = "" +phase = "fwd" + +kernel_profile_data = list(enumerate(reader)) + +for rid, row in kernel_profile_data: + if rid <= 2: + continue + kernel = row[4] + kid = rid - 2 + if "fused_classifier" in kernel: + # classifier: layernorm -> matmul -> fused -> bw matmul (x2) -> bw layernorm + CLS_START = kid - 2 + +assert CLS_START != -1 print() print("Kernel calls:") -for rid, row in enumerate(reader): +for rid, row in kernel_profile_data: if rid == 0: # headings - print(f"id pass {'name':<40} {'time':>8} {'RAM rd':>8} {'RAM wt':>8} {'L2 rd':>8} {'L2 wt':>8} {'inst':>8}") + print(f"id pass {'name':<40} {'time':>8} {'RAM rd':>8} {'RAM wt':>8} {'L2 rd':>8} {'L2 wt':>8} {'inst':>8}") continue if rid == 1: # units - units = f" {'':<40} {'ms':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'MInst':>8}" + units = f" {'':<40} {'ms':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'MInst':>8}" print(units) print("." * len(units)) continue if rid == 2: - CC = row[10] # actual data @@ -78,16 +91,27 @@ for rid, row in enumerate(reader): kid = rid - 2 - if kid == 0 or kid == ADAM_ID - 1: + multiplier = 1 + if "encoder" in kernel: pass_name = "enc" + if phase == "bwd": + phase = "bwd-enc" elif CLS_START <= kid < CLS_START + CLS_NUM: # the classifier part, counts only once pass_name = "cls" - elif kid == ADAM_ID: + phase = "bwd" + elif "adamw" in kernel: # encoder layer or adam pass_name = "opt" + # before the first optimizer run, we create weight copies. + # they aren't part of regular processing, so they get a multiplier + # of zero + elif phase == "bwd-enc": + pass_name = "init" + multiplier = 0 else: - pass_name = "fwd" if kid < CLS_START else "bwd" + pass_name = phase + multiplier = N_LAYERS time *= N_LAYERS read *= N_LAYERS write *= N_LAYERS @@ -99,9 +123,16 @@ for rid, row in enumerate(reader): # some names include the return value, others don't? if " " in fn_name: fn_name = fn_name.split(" ")[1] - if "cutlass" in fn_name: + if "<" in fn_name: fn_name = fn_name.split("<")[0] + + # group together matmul kernels + if "cutlass" in fn_name: pass + elif fn_name.startswith("ampere_bf16"): + fn_name = "ampere_bf16" + elif fn_name.startswith("cudnn_generated_fort_native_sdpa"): + fn_name = "cudnn_generated_fort_native_sdpa" else: no_cutlass += time @@ -110,29 +141,34 @@ for rid, row in enumerate(reader): l2_write = l2_write * 32 / 1024 / 1024 / 1024 summaries[fn_name] += time + counts[fn_name] += multiplier passes[pass_name] += time - total['time'] += time - total['read'] += read - total['write'] += write - total['l2_read'] += l2_read - total['l2_write'] += l2_write - total['inst'] += inst + if pass_name != "init": + total['time'] += time + total['read'] += read + total['write'] += write + total['l2_read'] += l2_read + total['l2_write'] += l2_write + total['inst'] += inst - print(f"{kid:02} {pass_name:4} {fn_name:<40} {time:8.2f} {read:8.2f} {write:8.2f} {l2_read:8.2f} {l2_write:8.2f} {inst:8.2f}") + pass_info = f"{pass_name}✗{multiplier}" + print(f"{kid:02} {pass_info:7} {fn_name:<40} {time:8.2f} {read:8.2f} {write:8.2f} {l2_read:8.2f} {l2_write:8.2f} {inst:8.2f}") total_time = total['time'] print("." * len(units)) -print(f" {'Total':<40} {total['time']:8.2f} {total['read']:8.2f} {total['write']:8.2f} {total['l2_read']:8.2f} {total['l2_write']:8.2f} {total['inst']:8.2f}") +print(f" {'Total':<40} {total['time']:8.2f} {total['read']:8.2f} {total['write']:8.2f} {total['l2_read']:8.2f} {total['l2_write']:8.2f} {total['inst']:8.2f}") print() print("Kernel type summaries:") -print(f" {'name':<40} {'time':>6} {'frac':>6}") -ordered = sorted(summaries.items(), key=lambda x: x[1], reverse=True) -for entry, value in ordered: +print(f" {'name':<40} {'time':>6} {'frac':>6} {'count':>6}") +ordered_time = sorted(summaries.items(), key=lambda x: x[1], reverse=True) +for entry, value in ordered_time: # crop entry to be at most 40 characters if len(entry) > 40: - entry = entry[:37] + "..." - print(f" {entry:<40} {value:6.2f} {100*value / total_time:6.2f}%") + entry_text = entry[:37] + "..." + else: + entry_text = entry + print(f" {entry_text:<40} {value:6.2f} {100*value / total_time:6.2f}% {counts[entry]:>6d}") ts = total_time / 1000 From 13eeeb20c8f8c53e590aac1a12dae45c54b1712e Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 3 May 2024 21:31:54 +0300 Subject: [PATCH 04/15] try to detect if we need sudo --- profile_gpt2cu.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index 448dd5a..8318443 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -17,10 +17,18 @@ if NCU is None: # build the executable subprocess.check_call(["make", "profile_gpt2cu", "NO_MULTI_GPU=1", "USE_CUDNN=1"]) +# try to see if profiling is allowed for non-root: +options = subprocess.check_output(["modprobe", "-c", "nvidia"], text=True) +can_profile = len([l for l in options.splitlines() if "NVreg_RestrictProfilingToAdminUsers=0" in l]) != 0 + # record metrics # --full and --import-source are entirely superfluous for this script, but you might want to # manually inspect `profile.ncu-rep`, so we keep it here -cmd = ["sudo", NCU, "--set", "full", "--import-source", "yes", "-o", "profile", "-f", "./profile_gpt2cu"] +cmd = [NCU, "--set", "full", "--import-source", "yes", "-o", "profile", "-f", "./profile_gpt2cu"] +# do we need to run under sudo +if not can_profile: + print("NVreg_RestrictProfilingToAdminUsers=1, running with sudo") + cmd = ["sudo"] + cmd subprocess.check_call(cmd) # generate csv From 7fabb12e2e101726d6553a4ca5a0cb04a3b9591c Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 3 May 2024 21:54:06 +0300 Subject: [PATCH 05/15] fixed instruction counts --- profile_gpt2cu.py | 1 + 1 file changed, 1 insertion(+) diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index 8318443..72ece6f 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -125,6 +125,7 @@ for rid, row in kernel_profile_data: write *= N_LAYERS l2_read *= N_LAYERS l2_write *= N_LAYERS + inst *= N_LAYERS # split at "(" -- argument list fn_name = kernel.split("(")[0] From 977d6a6c6fab61ba7078e1dd8b7126e313b6690e Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 3 May 2024 21:57:31 +0300 Subject: [PATCH 06/15] =?UTF-8?q?fix=20=C3=97=20symbol?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- profile_gpt2cu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/profile_gpt2cu.py b/profile_gpt2cu.py index 72ece6f..8e88473 100644 --- a/profile_gpt2cu.py +++ b/profile_gpt2cu.py @@ -160,7 +160,7 @@ for rid, row in kernel_profile_data: total['l2_write'] += l2_write total['inst'] += inst - pass_info = f"{pass_name}✗{multiplier}" + pass_info = f"{pass_name}×{multiplier}" print(f"{kid:02} {pass_info:7} {fn_name:<40} {time:8.2f} {read:8.2f} {write:8.2f} {l2_read:8.2f} {l2_write:8.2f} {inst:8.2f}") total_time = total['time'] From aa5dd2eeab035561ff06a3503df371f47b3b8e22 Mon Sep 17 00:00:00 2001 From: Andrej Date: Fri, 3 May 2024 12:43:08 -0700 Subject: [PATCH 07/15] Update README.md --- README.md | 45 ++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 0b63ee3..bcdfc37 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,12 @@ # llm.c -LLM training in simple, pure C/CUDA. There is no need for 245MB of PyTorch or 107MB of cPython. Training GPT-2 (CPU, fp32) is ~1,000 lines of clean code in the single file [train_gpt2.c](train_gpt2.c), and training it on GPU is ~2,000 lines (adds CUDA kernels) in [train_gpt2.cu](train_gpt2.cu). The code compiles and runs instantly, it exactly matches the PyTorch reference implementation, and it ~matches the speed of (compiled) PyTorch (fp32, no flash attention). I chose GPT-2 as the first working example because it is the grand-daddy of LLMs, the first time the modern stack was put together. +LLM training in simple, pure C/CUDA. There is no need for 245MB of PyTorch or 107MB of cPython. Training GPT-2 (CPU, fp32) is ~1,000 lines of clean code in the single file [train_gpt2.c](train_gpt2.c), and training it on GPU is ~3,000 lines (adds CUDA kernels) in [train_gpt2.cu](train_gpt2.cu). The code compiles and runs instantly, it exactly matches the PyTorch reference implementation, and currently slightly exceeds the speed of (compiled) PyTorch (with bf16, torch compile, and flash attention). I chose GPT-2 as the first working example because it is the grand-daddy of LLMs, the first time the modern stack was put together. -Our current goal is to reproduce GPT-2 with a multi-node, mixed-precision, efficient implementation. For an overview of current ongoing work, see the latest [State of the Union](https://github.com/karpathy/llm.c/discussions/224) post. +Our current goal is to reproduce GPT-2. For an overview of current ongoing work, see the latest [State of the Union](https://github.com/karpathy/llm.c/discussions/344) post. I'd like this repo to only maintain C and CUDA code. Ports of this repo to other languages are very welcome, but should be done in separate repos, and then I am happy to link to them below in the "notable forks" section, just like I did in [llama2.c notable forks](https://github.com/karpathy/llama2.c/tree/master?tab=readme-ov-file#notable-forks). -## quick start (GPU) +## quick start (GPU, slow but stable and educational) The "I don't care about anything I just want to train and I have a GPU" section. Run: @@ -20,6 +20,41 @@ make train_gpt2fp32cu The above lines (1) download the [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset, tokenize it with the GPT-2 Tokenizer, (2) download and save the GPT-2 (124M) weights, (3) init from them in C/CUDA and train for one epoch on tineshakespeare with AdamW (using batch size 4, context length 1024, total of 74 steps), evaluate validation loss, and sample some text. Note that in this quickstart we are using the fp32 version [train_gpt2_fp32.cu](train_gpt2_fp32.cu) of the CUDA code. Below in the CUDA section we document the current "mainline" [train_gpt2.cu](train_gpt2.cu), which is still being very actively developed, uses mixed precision, and runs ~2X faster. +## quick start (GPU, fast bleeding edge) + +I want to see it go fast. In this case switch to our mainline, most optimized `train_gpt2.cu` and also turn on flash attention. Run: + +```bash +pip install -r requirements.txt +python prepro_tinyshakespeare.py +python train_gpt2.py +make train_gpt2cu +./train_gpt2cu +``` + +If you additionally install cuDNN (see `Makefile` for instructions), you can also go faster with flash attention + +```bash +make train_gpt2cu USE_CUDNN=1 +./train_gpt2cu +``` + +Note that the default batch size is very low (4). If you have enough memory on your GPU, I recommend you increase this to e.g. 32: + +```bash +./train_gpt2cu -b 32 +``` + +My standard "prod" run with a nice GPU (e.g. A100 40GB) actually trains on TinyStories instead of TinyShakespeare, and looks like this: + +```bash +python prepro_tinystories.py +make train_gpt2cu USE_CUDNN=1 +./train_gpt2cu -i data/TinyStories -v 250 -s 250 -g 144 -o stories.log -b 32 +``` + +Where I decrease the frequency of validation loss and sampling to every 250 steps, sample 144 tokens during sampling stage (to fit ~one story), and at batch size 32. + ## quick start (CPU) The "I am so GPU poor that I don't even have one" section. No worries, run: @@ -213,7 +248,7 @@ make test_gpt2cu ./test_gpt2cu ``` -If you have the latest CUDA you should expect this to compile OK, and you should see ~2X improved speed (~1.86X to be precise). +If you have the latest CUDA you should expect this to compile OK, and you should see a lot more improved performance. **Flash Attention**. As of May 1, 2024 we now support 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: @@ -292,7 +327,7 @@ First, I want `llm.c` to be a place for education. E.g. our `dev/cuda` folder is 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. In that sense I'd be ok to only be at e.g. 90% of PyTorch speed, if it means we can remain at ~2,000 readable lines of code with minimal exotic dependencies. 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`. +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. From 67a82a611ef6135146e4bc496809d4d99d664f64 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Fri, 3 May 2024 23:08:44 +0300 Subject: [PATCH 08/15] first attempt at moving cudnn out of the main file for faster compiles --- Makefile | 15 ++- cudnn_att.cu | 319 ++++++++++++++++++++++++++++++++++++++++++++++++++ train_gpt2.cu | 280 ++------------------------------------------ 3 files changed, 339 insertions(+), 275 deletions(-) create mode 100644 cudnn_att.cu diff --git a/Makefile b/Makefile index 3af0225..865f6fa 100644 --- a/Makefile +++ b/Makefile @@ -194,20 +194,23 @@ train_gpt2: train_gpt2.c test_gpt2: test_gpt2.c $(CC) $(CFLAGS) $(INCLUDES) $(LDFLAGS) $< $(LDLIBS) $(OUTPUT_FILE) -train_gpt2cu: train_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) +cudnn_att.o: cudnn_att.cu + $(NVCC) -c $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) + +train_gpt2cu: train_gpt2.cu cudnn_att.o + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) cudnn_att.o train_gpt2fp32cu: train_gpt2_fp32.cu - $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) + $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) test_gpt2cu: test_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) + $(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) $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) + $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) profile_gpt2cu: profile_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(NVCC_LDFLAGS) $(CUDA_OUTPUT_FILE) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) clean: $(REMOVE_FILES) $(TARGETS) diff --git a/cudnn_att.cu b/cudnn_att.cu new file mode 100644 index 0000000..a9ce3bf --- /dev/null +++ b/cudnn_att.cu @@ -0,0 +1,319 @@ +// all cudnn-related functions are in this file, so that they don't need to be recompiled everytime +// we change some unrelated piece of the code. +// TODO this currently duplicates some of the utilities from the main file + +#include +#include +#include + +// Specific configurations based on the enabled precision +#if defined(ENABLE_FP32) +typedef float floatX; + +// use fp16 (note: this may require gradient scaler, currently not implemented!) +#elif defined(ENABLE_FP16) +typedef half floatX; +#define CUBLAS_LOWP CUDA_R_16F + +#else // Default to bfloat16 +typedef __nv_bfloat16 floatX; +#endif + +// CUDA error checking +static void cudaCheck(cudaError_t error, const char *file, int line) { + if (error != cudaSuccess) { + printf("[CUDA ERROR] at file %s:%d:\n%s\n", file, line, + cudaGetErrorString(error)); + exit(EXIT_FAILURE); + } +}; +#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__)) + +// Profiler utils +namespace { + class NvtxRange { + public: + NvtxRange(const char* s) { nvtxRangePush(s); } + + NvtxRange(const std::string& base_str, int number) { + std::string range_string = base_str + " " + std::to_string(number); + nvtxRangePush(range_string.c_str()); + } + + ~NvtxRange() { nvtxRangePop(); } + }; +} +#define NVTX_RANGE_FN() NvtxRange nvtx_range(__FUNCTION__) + +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 checkCudnnErr(err) assert((int)err == 0); + +using graph_tensors_fwd = std::tuple, +std::shared_ptr, // Q, +std::shared_ptr, // K, +std::shared_ptr, // V, +std::shared_ptr, // Attn_scale, +std::shared_ptr, // O +std::shared_ptr>; // Stats + +using graph_tensors_bwd = std::tuple, +std::shared_ptr, // Q, +std::shared_ptr, // K, +std::shared_ptr, // V, +std::shared_ptr, // O +std::shared_ptr, // dO +std::shared_ptr, // Stats +std::shared_ptr, // Attn_scale, +std::shared_ptr, // dQ, +std::shared_ptr, // dK, +std::shared_ptr>; // dV + +// Need a cache because graph->build_operation_graph() is slow but everything else seems fast +using cache_type_fwd = std::unordered_map; +using cache_type_bwd = std::unordered_map; + +// Loosely based on cuDNN frontend samples functions and massively simplified +template +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(); + 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; +} + +template +auto lookup_cache_or_build_graph_bwd(Args... args) { + static cache_type_bwd user_maintained_cache_bwd; + auto [B, NH, T, HS] = std::make_tuple(args...); + + auto graph = std::make_shared(); + graph->set_io_data_type(CUDNN_16BIT) + .set_intermediate_data_type(fe::DataType_t::FLOAT) + .set_compute_data_type(fe::DataType_t::FLOAT); + + // (B, N, 3, NH, HS) + // must come from inp (which means we also need to convert THAT to FP16) + auto Q = graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q") + .set_dim({B, NH, T, HS}) + .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); + auto K = graph->tensor(fe::graph::Tensor_attributes() + .set_name("K") + .set_dim({B, NH, T, HS}) + .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); + auto V = graph->tensor(fe::graph::Tensor_attributes() + .set_name("V") + .set_dim({B, NH, T, HS}) + .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); + auto O = graph->tensor(fe::graph::Tensor_attributes() + .set_name("O") + .set_dim({B, NH, T, HS}) + .set_stride({NH * HS * T, HS, NH * HS, 1})); + auto dO = graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO") + .set_dim({B, NH, T, HS}) + .set_stride({NH * HS * T, HS, NH * HS, 1})); + + auto stats = graph->tensor(fe::graph::Tensor_attributes() + .set_name("stats") + .set_dim({B, NH, T, 1}) + .set_stride({NH * T, T, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + 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_backward_options = fe::graph::SDPA_backward_attributes() + .set_name("flash_attention_backward") + .set_causal_mask(true) + .set_attn_scale(attn_scale); + + // Create the graph operation and get the output tensors back + auto [dQ, dK, dV] = graph->sdpa_backward(Q, K, V, O, dO, stats, sdpa_backward_options); + + dQ->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); + dK->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); + dV->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); + + assert(graph->validate().is_good()); + auto key = graph->key(); + auto it = user_maintained_cache_bwd.find(key); + if (it != user_maintained_cache_bwd.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, O, dO, stats, attn_scale, dQ, dK, dV); + user_maintained_cache_bwd.insert({key, tuple}); + return tuple; +} + +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 + int B, int T, int NH, int C) { + NVTX_RANGE_FN(); + int HS = C / NH; // number of features per head + bool is_inference_only = (stats == nullptr); + + // 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, 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()); +} + +void attention_backward_cudnn(floatX* dqkvr, // output + floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs + int B, int T, int NH, int C) { + NVTX_RANGE_FN(); + int HS = C / NH; // number of features per head + + // Get graph and tensors from cache (or generate it on first use) + auto [graph, Q, K, V, O, dO, Stats, attn_scale, dQ, dK, dV] = + lookup_cache_or_build_graph_bwd(B, NH, T, HS); + + // Prepare all the tensor pointers for executing the graph + void* devPtrQ = qkvr; + void* devPtrK = (qkvr + NH * HS); + void* devPtrV = (qkvr + 2 * NH * HS); + void* devPtrO = o; + void* devPtrdO = dout; + void* devPtrStats = stats; + float attn_scale_cpu = 1.0 / sqrtf(HS); + + void* devPtrdQ = dqkvr; + void* devPtrdK = (dqkvr + NH * HS); + void* devPtrdV = (dqkvr + 2 * NH * HS); + + // Build variant pack that links each tensor to its data pointer + std::unordered_map, void*> variant_pack = { + {Q, devPtrQ}, {K, devPtrK}, {V, devPtrV}, {O, devPtrO}, {dO, devPtrdO}, {Stats, devPtrStats}, + {dQ, devPtrdQ}, {dK, devPtrdK}, {dV, devPtrdV}, + {attn_scale, &attn_scale_cpu}}; + + // 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()); +} + +void create_cudnn() { + checkCudnnErr(cudnnCreate(&cudnn_handle)); +} + +void destroy_cudnn() { + if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } + checkCudnnErr(cudnnDestroy(cudnn_handle)); +} \ No newline at end of file diff --git a/train_gpt2.cu b/train_gpt2.cu index 084a7b8..3e90d71 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -115,21 +115,6 @@ const ncclDataType_t ncclFloatX = ncclBfloat16; #endif #endif -#ifdef ENABLE_CUDNN -#include -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 checkCudnnErr(err) assert((int)err == 0); -#endif // ENABLE_CUDNN - // ---------------------------------------------------------------------------- // CUDA utils @@ -495,258 +480,20 @@ void printf0(const char *format, ...) { // ---------------------------------------------------------------------------- // cuDNN path #ifdef ENABLE_CUDNN - -using graph_tensors_fwd = std::tuple, - std::shared_ptr, // Q, - std::shared_ptr, // K, - std::shared_ptr, // V, - std::shared_ptr, // Attn_scale, - std::shared_ptr, // O - std::shared_ptr>; // Stats - -using graph_tensors_bwd = std::tuple, - std::shared_ptr, // Q, - std::shared_ptr, // K, - std::shared_ptr, // V, - std::shared_ptr, // O - std::shared_ptr, // dO - std::shared_ptr, // Stats - std::shared_ptr, // Attn_scale, - std::shared_ptr, // dQ, - std::shared_ptr, // dK, - std::shared_ptr>; // dV - -// Need a cache because graph->build_operation_graph() is slow but everything else seems fast -using cache_type_fwd = std::unordered_map; -using cache_type_bwd = std::unordered_map; - -// Loosely based on cuDNN frontend samples functions and massively simplified -template -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(); - 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; -} - -template -auto lookup_cache_or_build_graph_bwd(Args... args) { - static cache_type_bwd user_maintained_cache_bwd; - auto [B, NH, T, HS] = std::make_tuple(args...); - - auto graph = std::make_shared(); - graph->set_io_data_type(CUDNN_16BIT) - .set_intermediate_data_type(fe::DataType_t::FLOAT) - .set_compute_data_type(fe::DataType_t::FLOAT); - - // (B, N, 3, NH, HS) - // must come from inp (which means we also need to convert THAT to FP16) - auto Q = graph->tensor(fe::graph::Tensor_attributes() - .set_name("Q") - .set_dim({B, NH, T, HS}) - .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); - auto K = graph->tensor(fe::graph::Tensor_attributes() - .set_name("K") - .set_dim({B, NH, T, HS}) - .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); - auto V = graph->tensor(fe::graph::Tensor_attributes() - .set_name("V") - .set_dim({B, NH, T, HS}) - .set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1})); - auto O = graph->tensor(fe::graph::Tensor_attributes() - .set_name("O") - .set_dim({B, NH, T, HS}) - .set_stride({NH * HS * T, HS, NH * HS, 1})); - auto dO = graph->tensor(fe::graph::Tensor_attributes() - .set_name("dO") - .set_dim({B, NH, T, HS}) - .set_stride({NH * HS * T, HS, NH * HS, 1})); - - auto stats = graph->tensor(fe::graph::Tensor_attributes() - .set_name("stats") - .set_dim({B, NH, T, 1}) - .set_stride({NH * T, T, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - 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_backward_options = fe::graph::SDPA_backward_attributes() - .set_name("flash_attention_backward") - .set_causal_mask(true) - .set_attn_scale(attn_scale); - - // Create the graph operation and get the output tensors back - auto [dQ, dK, dV] = graph->sdpa_backward(Q, K, V, O, dO, stats, sdpa_backward_options); - - dQ->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); - dK->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); - dV->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); - - assert(graph->validate().is_good()); - auto key = graph->key(); - auto it = user_maintained_cache_bwd.find(key); - if (it != user_maintained_cache_bwd.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, O, dO, stats, attn_scale, dQ, dK, dV); - user_maintained_cache_bwd.insert({key, tuple}); - return tuple; -} - +// functions defined in cudnn_att.cu +void create_cudnn(); +void destroy_cudnn(); 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 - int B, int T, int NH, int C) { - NVTX_RANGE_FN(); - int HS = C / NH; // number of features per head - bool is_inference_only = (stats == nullptr); - - // 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, 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()); -} + int B, int T, int NH, int C); void attention_backward_cudnn(floatX* dqkvr, // output floatX* dout, floatX* qkvr, floatX* o, float* stats, // inputs - int B, int T, int NH, int C) { - NVTX_RANGE_FN(); - int HS = C / NH; // number of features per head - - // Get graph and tensors from cache (or generate it on first use) - auto [graph, Q, K, V, O, dO, Stats, attn_scale, dQ, dK, dV] = - lookup_cache_or_build_graph_bwd(B, NH, T, HS); - - // Prepare all the tensor pointers for executing the graph - void* devPtrQ = qkvr; - void* devPtrK = (qkvr + NH * HS); - void* devPtrV = (qkvr + 2 * NH * HS); - void* devPtrO = o; - void* devPtrdO = dout; - void* devPtrStats = stats; - float attn_scale_cpu = 1.0 / sqrtf(HS); - - void* devPtrdQ = dqkvr; - void* devPtrdK = (dqkvr + NH * HS); - void* devPtrdV = (dqkvr + 2 * NH * HS); - - // Build variant pack that links each tensor to its data pointer - std::unordered_map, void*> variant_pack = { - {Q, devPtrQ}, {K, devPtrK}, {V, devPtrV}, {O, devPtrO}, {dO, devPtrdO}, {Stats, devPtrStats}, - {dQ, devPtrdQ}, {dK, devPtrdK}, {dV, devPtrdV}, - {attn_scale, &attn_scale_cpu}}; - - // 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()); -} + int B, int T, int NH, int C); +#else +void create_cudnn() {} +void destroy_cudnn() {} #endif // ENABLE_CUDNN // ---------------------------------------------------------------------------- @@ -2600,10 +2347,8 @@ 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 + // set up cuDNN (noop if not available) + create_cudnn(); printf0("| device | %-50s |\n", deviceProp.name); printf0("| TF32 | %-50s |\n", enable_tf32 ? "enabled" : "disabled"); @@ -2786,10 +2531,7 @@ 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)); } - checkCudnnErr(cudnnDestroy(cudnn_handle)); - #endif + destroy_cudnn(); cudaCheck(cudaFree(cublaslt_workspace)); cublasCheck(cublasDestroy(cublas_handle)); cublasCheck(cublasLtDestroy(cublaslt_handle)); From 03b7323d1eb819d2c03edca109c85106fb895e66 Mon Sep 17 00:00:00 2001 From: lancer Date: Fri, 3 May 2024 17:38:39 -0700 Subject: [PATCH 09/15] minor fix of the include --- train_gpt2.cu | 2 +- train_gpt2_fp32.cu | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 084a7b8..085b940 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -44,7 +44,7 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200), #include #include #include -#include + // GPU / CUDA related #include #include diff --git a/train_gpt2_fp32.cu b/train_gpt2_fp32.cu index ae07be6..178288f 100644 --- a/train_gpt2_fp32.cu +++ b/train_gpt2_fp32.cu @@ -13,14 +13,13 @@ the layernorms are connected to the residuals so we += in layernorm backward. #include #include -#include #include #include #include #include #include #include -#include + // GPU / CUDA related #include #include From aa5bb258b6a998b33f861f4ce2e949a06bd2cb3b Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 4 May 2024 11:59:59 +0300 Subject: [PATCH 10/15] fixed up test and profile targets --- Makefile | 8 ++++---- profile_gpt2.cu | 9 ++------- test_gpt2.cu | 11 +++-------- 3 files changed, 9 insertions(+), 19 deletions(-) diff --git a/Makefile b/Makefile index 865f6fa..8be333b 100644 --- a/Makefile +++ b/Makefile @@ -203,14 +203,14 @@ train_gpt2cu: train_gpt2.cu cudnn_att.o train_gpt2fp32cu: train_gpt2_fp32.cu $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) -test_gpt2cu: test_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) +test_gpt2cu: test_gpt2.cu cudnn_att.o + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) cudnn_att.o test_gpt2fp32cu: test_gpt2_fp32.cu $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) -profile_gpt2cu: profile_gpt2.cu - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) +profile_gpt2cu: profile_gpt2.cu cudnn_att.o + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) cudnn_att.o clean: $(REMOVE_FILES) $(TARGETS) diff --git a/profile_gpt2.cu b/profile_gpt2.cu index 5b043e0..2ea7f98 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -49,9 +49,7 @@ int main() { cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode)); // setup the (global) cuBLASLt workspace cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); - #ifdef ENABLE_CUDNN - checkCudnnErr(cudnnCreate(&cudnn_handle)); - #endif + create_cudnn(); // build the GPT-2 model from a checkpoint GPT2 model; @@ -81,10 +79,7 @@ int main() { // free gpt2_free(&model); - #ifdef ENABLE_CUDNN - if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } - checkCudnnErr(cudnnDestroy(cudnn_handle)); - #endif + destroy_cudnn(); cudaCheck(cudaFree(cublaslt_workspace)); cublasCheck(cublasDestroy(cublas_handle)); cublasCheck(cublasLtDestroy(cublaslt_handle)); diff --git a/test_gpt2.cu b/test_gpt2.cu index 67afa50..1a7df63 100644 --- a/test_gpt2.cu +++ b/test_gpt2.cu @@ -106,10 +106,8 @@ int main(int argc, char *argv[]) { cublasMath_t cublas_math_mode = enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH; cublasCheck(cublasSetMathMode(cublas_handle, cublas_math_mode)); cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size)); - - #ifdef ENABLE_CUDNN - checkCudnnErr(cudnnCreate(&cudnn_handle)); - #endif + // set up cuDNN (noop if not available) + create_cudnn(); // build the GPT-2 model from a checkpoint GPT2 model; @@ -326,10 +324,7 @@ int main(int argc, char *argv[]) { free(grads_memory_cpu); free(grads_memory_cpu_float); gpt2_free(&model); - #ifdef ENABLE_CUDNN - if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); } - checkCudnnErr(cudnnDestroy(cudnn_handle)); - #endif + destroy_cudnn(); cudaCheck(cudaFree(cublaslt_workspace)); cublasCheck(cublasDestroy(cublas_handle)); cublasCheck(cublasLtDestroy(cublaslt_handle)); From 19c290d7e603d7a5a9162665767d4249959efbfa Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sat, 4 May 2024 12:31:15 +0300 Subject: [PATCH 11/15] improved debugging for cudnn --- cudnn_att.cu | 62 ++++++++++++++++++++++++++++++---------------------- 1 file changed, 36 insertions(+), 26 deletions(-) diff --git a/cudnn_att.cu b/cudnn_att.cu index a9ce3bf..2735bbd 100644 --- a/cudnn_att.cu +++ b/cudnn_att.cu @@ -57,25 +57,35 @@ static size_t cudnn_workspace_size = 0; // dynamically allocated as needed (up t static void* cudnn_workspace = NULL; #define checkCudnnErr(err) assert((int)err == 0); +static void checkCudnnFE(fe::error_object e, const char *file, int line) { + if(!e.is_good()) { + printf("[CUDNN ERROR] at file %s:%d:\n%s\n", file, line, e.err_msg.c_str()); + exit(EXIT_FAILURE); + } +} +#define checkCudnnFE(err) checkCudnnFE(err, __FILE__, __LINE__) + using graph_tensors_fwd = std::tuple, -std::shared_ptr, // Q, -std::shared_ptr, // K, -std::shared_ptr, // V, -std::shared_ptr, // Attn_scale, -std::shared_ptr, // O -std::shared_ptr>; // Stats + std::shared_ptr, // Q, + std::shared_ptr, // K, + std::shared_ptr, // V, + std::shared_ptr, // Attn_scale, + std::shared_ptr, // O + std::shared_ptr // Stats +>; using graph_tensors_bwd = std::tuple, -std::shared_ptr, // Q, -std::shared_ptr, // K, -std::shared_ptr, // V, -std::shared_ptr, // O -std::shared_ptr, // dO -std::shared_ptr, // Stats -std::shared_ptr, // Attn_scale, -std::shared_ptr, // dQ, -std::shared_ptr, // dK, -std::shared_ptr>; // dV + std::shared_ptr, // Q, + std::shared_ptr, // K, + std::shared_ptr, // V, + std::shared_ptr, // O + std::shared_ptr, // dO + std::shared_ptr, // Stats + std::shared_ptr, // Attn_scale, + std::shared_ptr, // dQ, + std::shared_ptr, // dK, + std::shared_ptr // dV +>; // Need a cache because graph->build_operation_graph() is slow but everything else seems fast using cache_type_fwd = std::unordered_map; @@ -130,7 +140,7 @@ auto lookup_cache_or_build_graph_fwd(Args... args) { .set_stride({H * T, T, 1, 1}); } - assert(graph->validate().is_good()); + checkCudnnFE(graph->validate()); auto key = graph->key(); auto it = user_maintained_cache_fwd.find(key); if (it != user_maintained_cache_fwd.end()) { @@ -138,10 +148,10 @@ auto lookup_cache_or_build_graph_fwd(Args... args) { } // Build the operation graph and execution part (this is the VERY SLOW PART) - assert(graph->build_operation_graph(cudnn_handle).is_good()); + checkCudnnFE(graph->build_operation_graph(cudnn_handle)); 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()); + checkCudnnFE(graph->check_support(cudnn_handle)); + checkCudnnFE(graph->build_plans(cudnn_handle)); auto tuple = std::make_tuple(graph, Q, K, V, attn_scale, O, stats); user_maintained_cache_fwd.insert({key, tuple}); @@ -204,7 +214,7 @@ auto lookup_cache_or_build_graph_bwd(Args... args) { dK->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); dV->set_output(true).set_dim({B, NH, T, HS}).set_stride({3 * NH * HS * T, HS, 3 * NH * HS, 1}); - assert(graph->validate().is_good()); + checkCudnnFE(graph->validate()); auto key = graph->key(); auto it = user_maintained_cache_bwd.find(key); if (it != user_maintained_cache_bwd.end()) { @@ -212,10 +222,10 @@ auto lookup_cache_or_build_graph_bwd(Args... args) { } // Build the operation graph and execution part (this is the VERY SLOW PART) - assert(graph->build_operation_graph(cudnn_handle).is_good()); + checkCudnnFE(graph->build_operation_graph(cudnn_handle)); 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()); + checkCudnnFE(graph->check_support(cudnn_handle)); + checkCudnnFE(graph->build_plans(cudnn_handle)); auto tuple = std::make_tuple(graph, Q, K, V, O, dO, stats, attn_scale, dQ, dK, dV); user_maintained_cache_bwd.insert({key, tuple}); @@ -261,7 +271,7 @@ void attention_forward_cudnn(floatX* out, // output: (B, T, NH, HS) } // Execute graph - assert(graph->execute(cudnn_handle, variant_pack, cudnn_workspace).is_good()); + checkCudnnFE(graph->execute(cudnn_handle, variant_pack, cudnn_workspace)); cudaCheck(cudaGetLastError()); } @@ -305,7 +315,7 @@ void attention_backward_cudnn(floatX* dqkvr, } // Execute graph - assert(graph->execute(cudnn_handle, variant_pack, cudnn_workspace).is_good()); + checkCudnnFE(graph->execute(cudnn_handle, variant_pack, cudnn_workspace)); cudaCheck(cudaGetLastError()); } From 34304361bab643dbc80e24f200fb8aa257a659b7 Mon Sep 17 00:00:00 2001 From: ademeure Date: Sun, 5 May 2024 00:16:49 +0100 Subject: [PATCH 12/15] fix merge conflict --- profile_gpt2.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/profile_gpt2.cu b/profile_gpt2.cu index eab7fc5..cf22b71 100644 --- a/profile_gpt2.cu +++ b/profile_gpt2.cu @@ -71,8 +71,8 @@ int main() { GPT2 model; gpt2_build_from_checkpoint(&model, "gpt2_124M_bf16.bin"); - int B = 24; - int T = 1024; + int B = 24; // if program OOMs decrease this number, e.g. all the way down to 4 or etc + int T = 1024; // if even that OOMs move on to this one. keep them nice and powers of 2 printf("batch size: %d\n", B); printf("sequence length: %d\n", T); From b087b9c819516378cb0f60b00a35f9d812d94304 Mon Sep 17 00:00:00 2001 From: Erik Schultheis Date: Sun, 5 May 2024 02:44:15 +0300 Subject: [PATCH 13/15] don't compile/link cudnn if not asked for it --- Makefile | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index 8be333b..04cbfbb 100644 --- a/Makefile +++ b/Makefile @@ -18,6 +18,7 @@ NVCC_LDFLAGS = -lcublas -lcublasLt NVCC_INCLUDES = NVCC_LDLIBS = NCLL_INCUDES = +NVCC_CUDNN = # overridable flag for multi-GPU training. by default we won't build with cudnn # because it bloats up the compile time from a few seconds to ~minute USE_CUDNN ?= 0 @@ -81,6 +82,7 @@ ifeq ($(USE_CUDNN), 1) NVCC_INCLUDES += -I$(CUDNN_FRONTEND_PATH) NVCC_LDFLAGS += -lcudnn NVCC_FLAGS += -DENABLE_CUDNN + NVCC_CUDNN = cudnn_att.o else $(error ✗ cuDNN not found. See the Makefile for our currently hard-coded paths / install instructions) endif @@ -197,20 +199,20 @@ test_gpt2: test_gpt2.c cudnn_att.o: cudnn_att.cu $(NVCC) -c $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) -train_gpt2cu: train_gpt2.cu cudnn_att.o - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) cudnn_att.o +train_gpt2cu: train_gpt2.cu $(NVCC_CUDNN) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) $(NVCC_CUDNN) train_gpt2fp32cu: train_gpt2_fp32.cu $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) -test_gpt2cu: test_gpt2.cu cudnn_att.o - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) cudnn_att.o +test_gpt2cu: test_gpt2.cu $(NVCC_CUDNN) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) $(NVCC_CUDNN) test_gpt2fp32cu: test_gpt2_fp32.cu $(NVCC) $(NVCC_FLAGS) $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) -profile_gpt2cu: profile_gpt2.cu cudnn_att.o - $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) cudnn_att.o +profile_gpt2cu: profile_gpt2.cu $(NVCC_CUDNN) + $(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $< $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS) $(CUDA_OUTPUT_FILE) $(NVCC_CUDNN) clean: $(REMOVE_FILES) $(TARGETS) From a0b79358bdb9bcc005d67e519566736a9405b7e7 Mon Sep 17 00:00:00 2001 From: ademeure Date: Sun, 5 May 2024 01:01:59 +0100 Subject: [PATCH 14/15] Fix loss not being calculated for forward-only paths --- train_gpt2.cu | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index 971e9be..a59773f 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2059,7 +2059,7 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) { model->use_master_weights = 1; // keep master weights copy in float for optim update? } -void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { +void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, bool get_loss=true) { NVTX_RANGE_FN(); // targets are optional and could be NULL // in this function we must be careful and use size_t instead of int, otherwise @@ -2211,6 +2211,13 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) { // if we don't have targets, we don't have loss model->mean_loss = -1.0f; } + + // accumulate the loss immediately if we are not going to run gpt2_backward(), e.g. inference + if (get_loss) { + cudaCheck(cudaEventSynchronize(loss_event)); // hopefully finished long ago + for (int i=0; imean_loss += (float)(model->cpu_losses[i]); } + model->mean_loss /= B*T; + } } void gpt2_zero_grad(GPT2 *model) { @@ -2836,7 +2843,7 @@ int main(int argc, char *argv[]) { // if we're overfitting a single batch, we'll only call this at step = 0 dataloader_next_batch(&train_loader); } - gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T); + gpt2_forward(&model, train_loader.inputs, train_loader.targets, B, T, false); gpt2_zero_grad(&model); gpt2_backward(&model); if (multi_gpu_config.num_processes > 1) { From ce333de156edda1e3de7ad7b72c23ad2b5471168 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 5 May 2024 00:18:39 +0000 Subject: [PATCH 15/15] add ema to tok/s --- train_gpt2.cu | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/train_gpt2.cu b/train_gpt2.cu index f6cc3a2..e8fdbc7 100644 --- a/train_gpt2.cu +++ b/train_gpt2.cu @@ -2765,6 +2765,7 @@ int main(int argc, char *argv[]) { cudaCheck(cudaEventCreate(&end)); cudaCheck(cudaProfilerStart()); double total_sum_iteration_time_s = 0.0; + float ema_tokens_per_second = 0.0f; for (int step = 0; step <= train_num_batches; step++) { NvtxRange step_range("Train step", step); @@ -2857,12 +2858,18 @@ int main(int argc, char *argv[]) { cudaCheck(cudaEventSynchronize(end)); // wait for the end event to finish to get correct timings cudaCheck(cudaEventElapsedTime(&time_elapsed_ms, start, end)); + float tokens_per_second = multi_gpu_config.num_processes * (B * T) / time_elapsed_ms * 1000.0; + float bias_corrected_ema_tokens_per_second = tokens_per_second; // by default set to non-ema version if (step > 0) { // consider the first batch to be a warmup (e.g. cuBLAS/cuDNN initialisation) total_sum_iteration_time_s += time_elapsed_ms / 1000.0; + // smooth out the tok/s with an exponential moving average, and bias correct just like in AdamW + ema_tokens_per_second = 0.95f * ema_tokens_per_second + 0.05f * tokens_per_second; + bias_corrected_ema_tokens_per_second = ema_tokens_per_second / (1.0f - powf(0.95f, step)); } - int tokens_per_second = multi_gpu_config.num_processes * (B * T) / time_elapsed_ms * 1000.0; float accumulated_loss = multi_gpu_config.num_processes == 1 ? model.mean_loss : model.accumulated_mean_loss; - printf0("step %4d/%d: train loss %f (acc %f) (%f ms, %d tok/s)\n", step + 1, train_num_batches, model.mean_loss, accumulated_loss, time_elapsed_ms, tokens_per_second); + printf0("step %4d/%d: train loss %f (acc %f) (%f ms, %0f tok/s)\n", + step + 1, train_num_batches, model.mean_loss, accumulated_loss, + time_elapsed_ms, bias_corrected_ema_tokens_per_second); logger_log_train(&logger, step, model.mean_loss); // disable the profiler after 3 steps of optimization