aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge remote-tracking branch 'karpathy/master' into cleanup_may4

Browse source
This commit is contained in:
ademeure 2024-05-05 01:49:25 +01:00
commit c261eecda8
7 changed files with 502 additions and 300 deletions
Download patch file Download diff file Expand all files Collapse all files

21
Makefile
View file

@ -20,6 +20,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
@ -83,6 +84,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
@ -196,20 +198,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 $(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) $(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)
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) $(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) $(CUDA_OUTPUT_FILE)
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)

67
README.md
View file

@ -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,33 +327,37 @@ 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.
## notable forks
- C#
- [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): a C# port of this project
- 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 [CUDA MODE Discord Server](https://discord.gg/cudamode)
- Mojo
- [llm.🔥](https://github.com/dorjeduck/llm.mojo) by @[dorjeduck](https://github.com/dorjeduck): a Mojo 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
- 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
- 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

329
cudnn_att.cu Normal file
View file

@ -0,0 +1,329 @@
// 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 <cudnn_frontend.h>
#include <cuda_bf16.h>
#include <nvtx3/nvToolsExt.h>
// 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);
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<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
>;
using graph_tensors_bwd = 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>, // O
std::shared_ptr<fe::graph::Tensor_attributes>, // dO
std::shared_ptr<fe::graph::Tensor_attributes>, // Stats
std::shared_ptr<fe::graph::Tensor_attributes>, // Attn_scale,
std::shared_ptr<fe::graph::Tensor_attributes>, // dQ,
std::shared_ptr<fe::graph::Tensor_attributes>, // dK,
std::shared_ptr<fe::graph::Tensor_attributes> // dV
>;
// 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>;
using cache_type_bwd = std::unordered_map<std::size_t, graph_tensors_bwd>;
// 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});
}
checkCudnnFE(graph->validate());
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)
checkCudnnFE(graph->build_operation_graph(cudnn_handle));
auto plans = graph->create_execution_plans({fe::HeurMode_t::A});
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});
return tuple;
}
template <typename... Args>
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<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);
// (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});
checkCudnnFE(graph->validate());
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)
checkCudnnFE(graph->build_operation_graph(cudnn_handle));
auto plans = graph->create_execution_plans({fe::HeurMode_t::A});
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});
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<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
checkCudnnFE(graph->execute(cudnn_handle, variant_pack, cudnn_workspace));
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<std::shared_ptr<fe::graph::Tensor_attributes>, 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
checkCudnnFE(graph->execute(cudnn_handle, variant_pack, cudnn_workspace));
cudaCheck(cudaGetLastError());
}
void create_cudnn() {
checkCudnnErr(cudnnCreate(&cudnn_handle));
}
void destroy_cudnn() {
if (cudnn_workspace != NULL) { cudaCheck(cudaFree(cudnn_workspace)); }
checkCudnnErr(cudnnDestroy(cudnn_handle));
}

5
profile_gpt2.cu
View file

@ -34,8 +34,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);
@ -46,6 +46,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

96
profile_gpt2cu.py
View file

@ -14,13 +14,21 @@ 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"])
# 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 = [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
@ -39,32 +47,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,30 +99,49 @@ 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
l2_read *= N_LAYERS
l2_write *= N_LAYERS
inst *= N_LAYERS
# split at "(" -- argument list
fn_name = kernel.split("(")[0]
# 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,26 +150,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" {entry:<40} {value:6.2f} {100*value / total_time:6.2f}%")
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_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

281
train_gpt2.cu
View file

@ -1,5 +1,6 @@
/*
GPT-2 Transformer Neural Net trained in raw CUDA
GPT-2 Transformer Neural Net trained in raw CUDA
Non-trivial notes to be aware of:
We are being clever in the backward pass to conserve memory.
@ -92,19 +93,6 @@ const ncclDataType_t ncclFloatX = ncclBfloat16;
#endif
#endif
#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;
#endif // ENABLE_CUDNN
// ----------------------------------------------------------------------------
// CUDA utils
@ -439,237 +427,20 @@ void printf0(const char *format, ...) {
// ----------------------------------------------------------------------------
// cuDNN path
#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
using graph_tensors_bwd = 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>, // O
std::shared_ptr<fe::graph::Tensor_attributes>, // dO
std::shared_ptr<fe::graph::Tensor_attributes>, // Stats
std::shared_ptr<fe::graph::Tensor_attributes>, // Attn_scale,
std::shared_ptr<fe::graph::Tensor_attributes>, // dQ,
std::shared_ptr<fe::graph::Tensor_attributes>, // dK,
std::shared_ptr<fe::graph::Tensor_attributes>>; // dV
// 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>;
using cache_type_bwd = std::unordered_map<std::size_t, graph_tensors_bwd>;
// 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());
assert(graph->get_workspace_size() <= cudnn_workspace_size); // fwd shouldn't need workspace
auto tuple = std::make_tuple(graph, Q, K, V, attn_scale, O, stats);
user_maintained_cache_fwd.insert({key, tuple});
return tuple;
}
template <typename... Args>
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<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);
// (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());
// 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));
}
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<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;
}
// 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<std::shared_ptr<fe::graph::Tensor_attributes>, 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}};
// 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
// ----------------------------------------------------------------------------
@ -1912,7 +1683,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
@ -2064,6 +1835,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; i<B*T; i++) { model->mean_loss += (float)(model->cpu_losses[i]); }
model->mean_loss /= B*T;
}
}
void gpt2_zero_grad(GPT2 *model) {
@ -2316,16 +2094,15 @@ void common_start(bool override_enable_tf32 = true) {
cublasCheck(cublasSetStream(cublas_handle, main_stream));
cublasCheck(cublasLtCreate(&cublaslt_handle));
cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size));
#ifdef ENABLE_CUDNN
checkCudnnErr(cudnnCreate(&cudnn_handle));
#endif
// TF32 precision is equivalent to torch.set_float32_matmul_precision('high')
bool enable_tf32 = PRECISION_MODE == PRECISION_FP32 && deviceProp.major >= 8 && override_enable_tf32;
cublasCheck(cublasSetMathMode(cublas_handle, enable_tf32 ? CUBLAS_TF32_TENSOR_OP_MATH : CUBLAS_DEFAULT_MATH));
cublas_compute = enable_tf32 ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F;
// setup the (global) cuBLASLt workspace
cudaCheck(cudaMalloc(&cublaslt_workspace, cublaslt_workspace_size));
create_cudnn();
}
void common_free(GPT2 &model) {
@ -2338,13 +2115,10 @@ void common_free(GPT2 &model) {
cudaCheck(cudaStreamDestroy(main_stream));
gpt2_free(&model);
#ifdef ENABLE_CUDNN
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));
create_cudnn();
}
#ifndef TESTING
@ -2628,6 +2402,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);
@ -2706,7 +2481,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) {
@ -2720,12 +2495,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

3
train_gpt2_fp32.cu
View file

@ -13,14 +13,13 @@ the layernorms are connected to the residuals so we += in layernorm backward.
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <math.h>
#include <time.h>
#include <assert.h>
#include <float.h>
#include <string.h>
#include <unistd.h>
#include <assert.h>
// GPU / CUDA related
#include <cublas_v2.h>
#include <cuda_runtime.h>