separate out write_checkpoint and delete_checkpoint

Minor refactor

.github/workflows/ci_gpu.yml

@@ -39,13 +39,7 @@ jobs:
       - name: Train model (FP32) with gpt2_124M.bin
         run: |
           PRECISION=FP32 make train_gpt2cu
-          ./train_gpt2cu -b 1 -t 64 -d 256 -l 0.0001 -v 200 -s 200 -a 1 -x 10 -r 0 -f 0 -e "gpt2_124M.bin"
-
-      - name: Test for percent loss differential for FP32 
-        run: |
-          PRECISION=FP32 make train_gpt2cu
-          ./train_gpt2cu -b 1 -t 64 -d 256 -l 0.0001 -v 200 -s 200 -a 1 -x 10 -r 0 -f 0 -e "gpt2_124M.bin" > train_gpt2cu_fp32_precision.txt
-          python dev/loss_checker_ci.py -f train_gpt2cu_fp32_precision.txt -s 20 -e 28 -a 5.0
+          ./train_gpt2cu -b 4 -t 64 -l 1e-4 -v 200 -s 200 -a 1 -x 10 -e gpt2_124M.bin
 
       - name: Build FP32 precision
         run: PRECISION=FP32 make test_gpt2cu profile_gpt2cu

.github/workflows/ci_tests.yml

@@ -1,100 +0,0 @@
-name: Unit, Static and other Tests
-
-on:
-  create:
-  workflow_dispatch:
-  push:
-    branches:
-      - master
-  pull_request:
-    branches:
-      - master
-
-jobs:
-  dataloader_test:
-    runs-on: ubuntu-latest
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: test the dataloader without / with sanitize address
-        run: |
-          cd dev/test
-          make PRECISION=BF16 test_dataloader
-          ./test_dataloader   
-          make clean       
-          make PRECISION=BF16 TEST_CFLAGS="-fsanitize=address -fno-omit-frame-pointer" test_dataloader 
-          ./test_dataloader          
-
-  ptx_and_sass_files:
-    runs-on: ubuntu-latest
-    container:
-      image: nvidia/cuda:12.4.1-devel-ubuntu22.04
-
-    steps:
-      - name: Checkout code
-        uses: actions/checkout@v4
-
-      - name: Install OpenMP and OpenMPI
-        run: apt-get update && apt-get install -y libomp-dev libopenmpi-dev
-    
-      - name: Generate ptx/sass files and upload them to persistent storage
-        run: |
-          mkdir -p dev/cuda/ptx_sass_logs
-          make train_gpt2cu
-          cuobjdump --dump-ptx train_gpt2cu > dev/cuda/train_gpt2cu.ptx
-          cuobjdump --dump-sass train_gpt2cu > dev/cuda/train_gpt2cu.sass          
-          cd dev/cuda
-          make -j all_ptx
-          make -j all_sass
-          cp *.ptx ptx_sass_logs/
-          cp *.sass ptx_sass_logs/
-          ls ptx_sass_logs/
-
-      - name: Generate ptx/sass files for A100 and upload them to persistent storage
-        run: |
-            mkdir -p dev/cuda/ptx_sass_logs_A100
-            make train_gpt2cu GPU_COMPUTE_CAPABILITY=80
-            cuobjdump --dump-ptx train_gpt2cu > dev/cuda/train_gpt2cu.ptx
-            cuobjdump --dump-sass train_gpt2cu > dev/cuda/train_gpt2cu.sass          
-            cd dev/cuda
-            make -j GPU_COMPUTE_CAPABILITY=80 all_ptx
-            make -j GPU_COMPUTE_CAPABILITY=80 all_sass
-            cp *.ptx ptx_sass_logs_A100/
-            cp *.sass ptx_sass_logs_A100/
-            ls ptx_sass_logs_A100/
-
-      - name: Generate ptx/sass files for H100 and upload them to persistent storage
-        run: |
-            mkdir -p dev/cuda/ptx_sass_logs_H100
-            make train_gpt2cu GPU_COMPUTE_CAPABILITY=90
-            cuobjdump --dump-ptx train_gpt2cu > dev/cuda/train_gpt2cu.ptx
-            cuobjdump --dump-sass train_gpt2cu > dev/cuda/train_gpt2cu.sass          
-            cd dev/cuda
-            make -j GPU_COMPUTE_CAPABILITY=90 all_ptx
-            make -j GPU_COMPUTE_CAPABILITY=90 all_sass
-            cp *.ptx ptx_sass_logs_H100/
-            cp *.sass ptx_sass_logs_H100/
-            ls ptx_sass_logs_H100/
-
-      - name: Upload ptx/sass files
-        uses: actions/upload-artifact@v4
-        with:
-          name: ptx_sass_files
-          path: dev/cuda/ptx_sass_logs/
-          retention-days: 30 # days to retain
-
-      - name: Upload ptx/sass files for A100
-        uses: actions/upload-artifact@v4
-        with:
-          name: ptx_sass_files_A100
-          path: dev/cuda/ptx_sass_logs_A100/
-          retention-days: 30 # days to retain          
-
-      - name: Upload ptx/sass files for H100
-        uses: actions/upload-artifact@v4
-        with:
-          name: ptx_sass_files_H100
-          path: dev/cuda/ptx_sass_logs_H100/
-          retention-days: 30 # days to retain                    

Makefile

@@ -11,12 +11,9 @@ REMOVE_FILES = rm -f
 OUTPUT_FILE = -o $@
 CUDA_OUTPUT_FILE = -o $@
 
-# Default O3 CPU optimization level for NVCC (0 for fastest compile time)
-FORCE_NVCC_O ?= 3
-
 # NVCC flags
 # -t=0 is short for --threads, 0 = number of CPUs on the machine
-NVCC_FLAGS = --threads=0 -t=0 --use_fast_math -std=c++17 -O$(FORCE_NVCC_O)
+NVCC_FLAGS = -O3 -t=0 --use_fast_math -std=c++17
 NVCC_LDFLAGS = -lcublas -lcublasLt
 NVCC_INCLUDES =
 NVCC_LDLIBS =
@@ -48,10 +45,8 @@ endif
 
 ifneq ($(CI),true) # if not in CI, then use the GPU query
   ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY=
-    ifneq ($(call file_exists_in_path, nvidia-smi),)
-      # Get the compute capabilities of all GPUs
-      # Remove decimal points, sort numerically in ascending order, and select the first (lowest) value
-      GPU_COMPUTE_CAPABILITY=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed 's/\.//g' | sort -n | head -n 1)
+    ifneq ($(call file_exists_in_path, __nvcc_device_query),)
+      GPU_COMPUTE_CAPABILITY = $(shell __nvcc_device_query)
       GPU_COMPUTE_CAPABILITY := $(strip $(GPU_COMPUTE_CAPABILITY))
     endif
   endif
@@ -67,7 +62,6 @@ $(info ---------------------------------------------)
 
 ifneq ($(OS), Windows_NT)
   NVCC := $(shell which nvcc 2>/dev/null)
-  NVCC_LDFLAGS += -lnvidia-ml
 
   # Function to test if the compiler accepts a given flag.
   define check_and_add_flag
@@ -204,7 +198,6 @@ else
       $(info ✗ Multi-GPU on CUDA on Darwin is not supported, skipping NCCL support)
     else ifeq ($(shell dpkg -l | grep -q nccl && echo "exists"), exists)
       $(info ✓ NCCL found, OK to train with multiple GPUs)
-      NVCC_FLAGS += -DMULTI_GPU
       NVCC_LDLIBS += -lnccl
     else
       $(info ✗ NCCL is not found, disabling multi-GPU support)
@@ -213,18 +206,15 @@ else
   endif
 endif
 
-# Attempt to find and include OpenMPI on the system
-OPENMPI_DIR ?= /usr/lib/x86_64-linux-gnu/openmpi
-OPENMPI_LIB_PATH = $(OPENMPI_DIR)/lib/
-OPENMPI_INCLUDE_PATH = $(OPENMPI_DIR)/include/
 ifeq ($(NO_USE_MPI), 1)
   $(info → MPI is manually disabled)
-else ifeq ($(shell [ -d $(OPENMPI_LIB_PATH) ] && [ -d $(OPENMPI_INCLUDE_PATH) ] && echo "exists"), exists)
+else ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists)
   $(info ✓ MPI enabled)
-  NVCC_INCLUDES += -I$(OPENMPI_INCLUDE_PATH)
-  NVCC_LDFLAGS += -L$(OPENMPI_LIB_PATH)
+  NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include
+  NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/
   NVCC_LDLIBS += -lmpi
   NVCC_FLAGS += -DUSE_MPI
+  NVCC_FLAGS += -DMULTI_GPU
 else
   $(info ✗ MPI not found)
 endif
@@ -286,5 +276,5 @@ profile_gpt2cu: profile_gpt2.cu $(NVCC_CUDNN)
 	$(NVCC) $(NVCC_FLAGS) $(PFLAGS) -lineinfo $^ $(NVCC_LDFLAGS) $(NVCC_INCLUDES) $(NVCC_LDLIBS)  $(CUDA_OUTPUT_FILE)
 
 clean:
-	$(REMOVE_FILES) $(TARGETS)
+	$(REMOVE_FILES) $(TARGETS) 
 	$(REMOVE_BUILD_OBJECT_FILES)

README.en.md

@@ -1,267 +0,0 @@
-# llm.c
-
-LLMs in simple, pure C/CUDA with no need for 245MB of PyTorch or 107MB of cPython. Current focus is on pretraining, in particular reproducing the [GPT-2](https://github.com/openai/gpt-2) and [GPT-3](https://arxiv.org/abs/2005.14165) miniseries, along with a parallel PyTorch reference implementation in [train_gpt2.py](train_gpt2.py). You'll recognize this file as a slightly tweaked [nanoGPT](https://github.com/karpathy/nanoGPT), an earlier project of mine. Currently, llm.c is a bit faster than PyTorch Nightly (by about 7%). In addition to the bleeding edge mainline code in [train_gpt2.cu](train_gpt2.cu), we have a simple reference CPU fp32 implementation in ~1,000 lines of clean code in one file [train_gpt2.c](train_gpt2.c). I'd like this repo to only maintain C and CUDA code. Ports to other languages or repos are very welcome, but should be done in separate repos, and I am happy to link to them below in the "notable forks" section. Developer coordination happens in the [Discussions](https://github.com/karpathy/llm.c/discussions) and on Discord, either the `#llmc` channel on the [Zero to Hero](https://discord.gg/3zy8kqD9Cp) channel, or on `#llmdotc` on [GPU MODE](https://discord.gg/gpumode) Discord.
-
-## quick start
-
-The best introduction to the llm.c repo today is reproducing the GPT-2 (124M) model. [Discussion #481](https://github.com/karpathy/llm.c/discussions/481) steps through this in detail. We can reproduce other models from the GPT-2 and GPT-3 series in both llm.c and in the parallel implementation of PyTorch. Have a look at the [scripts README](scripts/README.md).
-
-debugging tip: when you run the `make` command to build the binary, modify it by replacing `-O3` with `-g` so you can step through the code in your favorite IDE (e.g. vscode).
-
-## quick start (1 GPU, fp32 only)
-
-If you won't be training on multiple nodes, aren't interested in mixed precision, and are interested in learning CUDA, the fp32 (legacy) files might be of interest to you. These are files that were "checkpointed" early in the history of llm.c and frozen in time. They are simpler, more portable, and possibly easier to understand. Run the 1 GPU, fp32 code like this:
-
-```bash
-chmod u+x ./dev/download_starter_pack.sh
-./dev/download_starter_pack.sh
-make train_gpt2fp32cu
-./train_gpt2fp32cu
-```
-
-The download_starter_pack.sh script is a quick & easy way to get started and it downloads a bunch of .bin files that help get you off the ground. These contain: 1) the GPT-2 124M model saved in fp32, in bfloat16, 2) a "debug state" used in unit testing (a small batch of data, and target activations and gradients), 3) the GPT-2 tokenizer, and 3) the tokenized [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset. Alternatively, instead of running the .sh script, you can re-create these artifacts manually as follows:
-
-```bash
-pip install -r requirements.txt
-python dev/data/tinyshakespeare.py
-python train_gpt2.py
-```
-
-## quick start (CPU)
-
-The "I am so GPU poor that I don't even have one GPU" section. You can still enjoy seeing llm.c train! But you won't go too far. Just like the fp32 version above, the CPU version is an even earlier checkpoint in the history of llm.c, back when it was just a simple reference implementation in C. For example, instead of training from scratch, you can finetune a GPT-2 small (124M) to output Shakespeare-like text, as an example:
-
-```bash
-chmod u+x ./dev/download_starter_pack.sh
-./dev/download_starter_pack.sh
-make train_gpt2
-OMP_NUM_THREADS=8 ./train_gpt2
-```
-
-If you'd prefer to avoid running the starter pack script, then as mentioned in the previous section you can reproduce the exact same .bin files and artifacts by running `python dev/data/tinyshakespeare.py` and then `python train_gpt2.py`.
-
-The above lines (1) download an already tokenized [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset and download the GPT-2 (124M) weights, (3) init from them in C and train for 40 steps on tineshakespeare with AdamW (using batch size 4, context length only 64), evaluate validation loss, and sample some text. Honestly, unless you have a beefy CPU (and can crank up the number of OMP threads in the launch command), you're not going to get that far on CPU training LLMs, but it might be a good demo/reference. The output looks like this on my MacBook Pro (Apple Silicon M3 Max):
-
-```
-[GPT-2]
-max_seq_len: 1024
-vocab_size: 50257
-num_layers: 12
-num_heads: 12
-channels: 768
-num_parameters: 124439808
-train dataset num_batches: 1192
-val dataset num_batches: 128
-num_activations: 73323776
-val loss 5.252026
-step 0: train loss 5.356189 (took 1452.121000 ms)
-step 1: train loss 4.301069 (took 1288.673000 ms)
-step 2: train loss 4.623322 (took 1369.394000 ms)
-step 3: train loss 4.600470 (took 1290.761000 ms)
-... (trunctated) ...
-step 39: train loss 3.970751 (took 1323.779000 ms)
-val loss 4.107781
-generating:
----
-Come Running Away,
-Greater conquer
-With the Imperial blood
-the heaviest host of the gods
-into this wondrous world beyond.
-I will not back thee, for how sweet after birth
-Netflix against repounder,
-will not
-flourish against the earlocks of
-Allay
----
-```
-
-## datasets
-
-The data files inside `/dev/data/(dataset).py` are responsible for downloading, tokenizing and saving the tokens to .bin files, readable easily from C. So for example when you run:
-
-```bash
-python dev/data/tinyshakespeare.py
-```
-
-We download and tokenize the [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset. The output of this looks like this:
-
-```
-writing 32,768 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_val.bin
-writing 305,260 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_train.bin
-```
-
-The .bin files contain a short header (1024 bytes) and then a stream of tokens in uint16, indicating the token ids with the GPT-2 tokenizer. More datasets are available in `/dev/data`.
-
-## test
-
-I am also attaching a simple unit test for making sure our C code agrees with the PyTorch code. On the CPU as an example, compile and run with:
-
-```bash
-make test_gpt2
-./test_gpt2
-```
-
-This now loads the `gpt2_124M_debug_state.bin` file that gets written by train_gpt2.py, runs a forward pass, compares the logits and loss with the PyTorch reference implementation, then it does 10 iterations of training with Adam and makes sure the losses match PyTorch. To test the GPU version we run:
-
-```bash
-# fp32 test (cudnn not supported)
-make test_gpt2cu PRECISION=FP32 && ./test_gpt2cu
-# mixed precision cudnn test
-make test_gpt2cu USE_CUDNN=1 && ./test_gpt2cu
-```
-
-This tests both the fp32 path and the mixed precision path. The test should pass and print `overall okay: 1`.
-
-## tutorial
-
-I attached a very small tutorial here, in [doc/layernorm/layernorm.md](doc/layernorm/layernorm.md). It's a simple, step-by-step guide to implementing a single layer of the GPT-2 model, the layernorm layer. This is a good starting point to understand how the layers are implemented in C.
-
-**flash attention**. As of May 1, 2024 we use the Flash Attention from cuDNN. Because cuDNN bloats the compile time from a few seconds to ~minute and this code path is right now very new, this is disabled by default. You can enable it by compiling like this:
-
-```bash
-make train_gpt2cu USE_CUDNN=1
-```
-
-This will try to compile with cudnn and run it. You have to have cuDNN installed on your system. The [cuDNN installation instructions](https://developer.nvidia.com/cudnn) with apt-get will grab the default set of cuDNN packages. For a minimal setup, the cuDNN dev package is sufficient, e.g. on Ubuntu 22.04 for CUDA 12.x:
-
-```bash
-wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
-sudo dpkg -i cuda-keyring_1.1-1_all.deb
-sudo apt-get update
-sudo apt-get -y install libcudnn9-dev-cuda-12
-```
-
-On top of this you need the [cuDNN frontend](https://github.com/NVIDIA/cudnn-frontend/tree/main), but this is just header files. Simply clone the repo to your disk. The Makefile currently looks for it in either your home directory or the current directory. If you have put it elsewhere, add `CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include` to the `make` command-line.
-
-## multi-GPU training
-
-Make sure you install MPI and NCCL, e.g. on Linux:
-
-```bash
-sudo apt install openmpi-bin openmpi-doc libopenmpi-dev
-```
-
-For NCCL follow the instructions from the [official website](https://developer.nvidia.com/nccl/nccl-download) (e.g. network installer)
-
-and then:
-
-```bash
-make train_gpt2cu
-mpirun -np <number of GPUs> ./train_gpt2cu
-```
-
-or simply run one of our scripts under `./scripts/`.
-
-## multi-node training
-
-Make sure you've installed `NCCL` following instructions from [multi-GPU](#multi-gpu-training) section.
-
-There are 3 ways we currently support that allow you to run multi-node training:
-1) Use OpenMPI to exchange nccl id and initialize NCCL. See e.g. `./scripts/multi_node/run_gpt2_124M_mpi.sh` script for details.
-2) Use shared file system to init NCCL. See `./scripts/multi_node/run_gpt2_124M_fs.sbatch` script for details.
-3) Use TCP sockets to init NCCL. See `./scripts/multi_node/run_gpt2_124M_tcp.sbatch` script for details.
-
-Note:
-* If you're running in a slurm environment and your slurm doesn't support PMIx (which we assume will be a common situation given that `slurm-wlm` dropped PMIx support) you will have to use FS (2) or TCP (3) approach. To test whether your slurm supports PMIx run: `srun --mpi=list` and see whether you get `pmix` in the output.
-* If you don't have slurm set up, you can kick off a multi-node run using `mpirun` - MPI (1).
-
-None of these 3 methods is superior, we just offer you options so that you can run in your specific environment.
-
-## experiments / sweeps
-
-Just as an example process to sweep learning rates on a machine with 4 GPUs on TinyStories. Run a shell script `sweep.sh` (after you of course `chmod u+x sweep.sh`):
-
-```bash
-#!/bin/bash
-
-learning_rates=(3e-5 1e-4 3e-4 1e-3)
-
-for i in {0..3}; do
-    export CUDA_VISIBLE_DEVICES=$i
-    screen -dmS "tr$i" bash -c "./train_gpt2cu -i data/TinyStories -v 250 -s 250 -g 144 -l ${learning_rates[$i]} -o stories$i.log"
-done
-
-# you can bring these down with
-# screen -ls | grep -E "tr[0-3]" | cut -d. -f1 | xargs -I {} screen -X -S {} quit
-```
-
-This example opens up 4 screen sessions and runs the four commands with different LRs. This writes the log files `stories$i.log` with all the losses, which you can plot as you wish in Python. A quick example of how to parse and plot these logfiles is in [dev/vislog.ipynb](dev/vislog.ipynb).
-
-## repo
-
-A few more words on what I want this repo to be:
-
-First, I want `llm.c` to be a place for education. E.g. our `dev/cuda` folder is a place for a library of kernels for all the layers that are manually hand-written and very well documented, starting from very simple kernels all the way to more complex / faster kernels. If you have a new kernel with various different tradeoffs, please feel free to contribute it here.
-
-That said, I also want `llm.c` to be very fast too, even practically useful to train networks. E.g. to start, we should be able to reproduce the big GPT-2 (1.6B) training run. This requires that we incorporate whatever fastest kernels there are, including the use of libraries such as cuBLAS, cuBLASLt, CUTLASS, cuDNN, etc. I also think doing so serves an educational purpose to establish an expert upper bound, and a unit of measurement, e.g. you could say that your manually written kernels are 80% of cuBLAS speed, etc. Then you can choose to do a super fast run, or you can choose to "drag and drop" whatever manual kernels you wish to use, and run with those.
-
-However, as a constraint, I want to keep the mainline `llm.c` in the root folder simple and readable. If there is a PR that e.g. improves performance by 2% but it "costs" 500 lines of complex C code, and maybe an exotic 3rd party dependency, I may reject the PR because the complexity is not worth it. As a concrete example - making cuBLAS for matmuls the default in the root training loop is a no-brainer: it makes the mainline code much faster, it is a single line of interpretable code, and it is a very common dependency. On the side of this, we can have manual implementations that can compete with cuBLAS in `dev/cuda`.
-
-Lastly, I will be a lot more sensitive to complexity in the root folder of the project, which contains the main / default files of the project. In comparison, the `dev/` folder is a bit more of a scratch space for us to develop a library of kernels or classes and share useful or related or educational code, and some of this code could be ok to be (locally) complex.
-
-## notable forks
-
-- AMD support
-  - [llm.c](https://github.com/anthonix/llm.c) by @[anthonix](https://github.com/anthonix): support for AMD devices, such as the 7900 XTX
-
-- C#
-  - [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): a C# port of this project
-  - [Llm.cs](https://github.com/nietras/Llm.cs) by @[nietras](https://github.com/nietras): a C# port of this project with focus on easy to get started on any platform. Clone and run ✅
-
-- CUDA C++
-  - [llm.cpp](https://github.com/gevtushenko/llm.c) by @[gevtushenko](https://github.com/gevtushenko): a port of this project using the [CUDA C++ Core Libraries](https://github.com/NVIDIA/cccl)
-     - A presentation this fork was covered in [this lecture](https://www.youtube.com/watch?v=WiB_3Csfj_Q) in the [GPU MODE Discord Server](https://discord.gg/cudamode)
-
-- C++/CUDA
-  - [llm.cpp](https://github.com/zhangpiu/llm.cpp/tree/master/llmcpp) by @[zhangpiu](https://github.com/zhangpiu): a port of this project using the [Eigen](https://gitlab.com/libeigen/eigen), supporting CPU/CUDA.
-
-- WebGPU C++
-  - [gpu.cpp](https://github.com/AnswerDotAI/gpu.cpp) by @[austinvhuang](https://github.com/austinvhuang): a library for portable GPU compute in C++ using native WebGPU. Aims to be a general-purpose library, but also porting llm.c kernels to WGSL.
-  
-- C++
-  - [llm.cpp](https://github.com/GaoYusong/llm.cpp) by @[GaoYusong](https://github.com/GaoYusong): a port of this project featuring a C++ single-header [tinytorch.hpp](https://github.com/GaoYusong/llm.cpp/blob/main/tinytorch.hpp) library
-
-- Go
-  - [llm.go](https://github.com/joshcarp/llm.go) by @[joshcarp](https://github.com/joshcarp): a Go port of this project
-
-- Java
-  - [llm.java](https://github.com/harryjackson/llm.java) by @[harryjackson](https://github.com/harryjackson): a Java port of this project
-
-- Metal
-  - [llm.metal](https://github.com/regrettable-username/llm.metal) by @[regrettable-username](https://github.com/regrettable-username): LLM training in simple, raw C/Metal Shading Language
-
-- Mojo
-  - [llm.🔥](https://github.com/dorjeduck/llm.mojo) by @[dorjeduck](https://github.com/dorjeduck): a Mojo port of this project
-
-- OpenCL
-  - [llm.c](https://github.com/krrishnarraj/llm.c) by @[krrishnarraj](https://github.com/krrishnarraj): an OpenCL port of this project
-
-- Rust
-  -  [llm.rs](https://github.com/yijunyu/llm.rs) by @[Yijun Yu](https://github.com/yijunyu): a Rust rewrite with the aim to have same performance
-  -  [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project
-
-- Swift
-  - [llm.swift](https://github.com/otabuzzman/llm.swift) by @[otabuzzman](https://github.com/otabuzzman): a Swift port of this project
-
-- Zig
-  - [llm.zig](https://github.com/Saimirbaci/llm.zig) by @[saimirbaci](https://github.com/Saimirbaci): a Zig port of this project
- 
-- Habana Gaudi2
-  - [llm.tpc](https://github.com/abhilash1910/llm.tpc) by @[abhilash1910](https://github.com/abhilash1910): a Habana Gaudi2 port of this project 
-
-- Nim
-  - [llm.nim](https://github.com/Vindaar/llm.nim) by @[Vindaar](https://github.com/Vindaar): a Nim port of this project
-
-## discussions
-
-Ways of organizing development:
-
-- Experiencing a concrete issue with the repo? Use [Issues](https://github.com/karpathy/llm.c/issues).
-- Have some code to contribute? Open a [PR](https://github.com/karpathy/llm.c/pulls)
-- Chat about the repo, ask questions, etc.? Look at [Discussions](https://github.com/karpathy/llm.c/discussions).
-- Something faster? I created a new `#llmc` channel on my [Zero to Hero Discord channel](https://discord.gg/3zy8kqD9Cp).
-
-## license
-
-MIT

README.md

@@ -1,22 +1,16 @@
-<!-- WEHUB_ZH_README -->
-> [!NOTE]
-> 本文档由 WeHub 基于上游 README 翻译整理，属于社区翻译，非官方中文文档。  
-> [English](./README.en.md) · [原始项目](https://github.com/karpathy/llm.c) · [上游 README](https://github.com/karpathy/llm.c/blob/HEAD/README.md)  
-> 原作者、版权与许可证归属以原始项目及本仓库 LICENSE 文件为准。
-
 # llm.c
 
-用简洁、纯粹的 C/CUDA 实现 LLM，无需 245MB 的 PyTorch 或 107MB 的 cPython。当前重点是预训练（pretraining），尤其是复现 [GPT-2](https://github.com/openai/gpt-2) 和 [GPT-3](https://arxiv.org/abs/2005.14165) 迷你系列，同时在 [train_gpt2.py](train_gpt2.py) 中提供并行的 PyTorch 参考实现。你会认出该文件是我早期项目 [nanoGPT](https://github.com/karpathy/nanoGPT), 的稍加改动版本。目前，llm.c 比 PyTorch Nightly 略快（约 7%）。除了 [train_gpt2.cu](train_gpt2.cu) 中的前沿主线代码外，我们还在单文件 [train_gpt2.c](train_gpt2.c) 中提供了一个约 1,000 行整洁代码的简单 CPU fp32 参考实现。我希望本仓库仅维护 C 和 CUDA 代码。欢迎移植到其他语言或仓库，但应在独立仓库中进行，我很乐意在下方「notable forks」部分链接它们。开发者协作在 [Discussions](https://github.com/karpathy/llm.c/discussions) 以及 Discord 上进行，可在 [Zero to Hero](https://discord.gg/3zy8kqD9Cp) 频道的 `#llmc` 频道，或在 [GPU MODE](https://discord.gg/gpumode) Discord 的 `#llmdotc` 上交流。
+LLMs in simple, pure C/CUDA with no need for 245MB of PyTorch or 107MB of cPython. Current focus is on pretraining, in particular reproducing the [GPT-2](https://github.com/openai/gpt-2) and [GPT-3](https://arxiv.org/abs/2005.14165) miniseries, along with a parallel PyTorch reference implementation in [train_gpt2.py](train_gpt2.py). You'll recognize this file as a slightly tweaked [nanoGPT](https://github.com/karpathy/nanoGPT), an earlier project of mine. Currently, llm.c is a bit faster than PyTorch Nightly (by about 7%). In addition to the bleeding edge mainline code in [train_gpt2.cu](train_gpt2.cu), we have a simple reference CPU fp32 implementation in ~1,000 lines of clean code in one file [train_gpt2.c](train_gpt2.c). I'd like this repo to only maintain C and CUDA code. Ports to other languages or repos are very welcome, but should be done in separate repos, and I am happy to link to them below in the "notable forks" section. Developer coordination happens in the [Discussions](https://github.com/karpathy/llm.c/discussions) and on Discord, either the `#llmc` channel on the [Zero to Hero](https://discord.gg/3zy8kqD9Cp) channel, or on `#llmdotc` on CUDA MODE Discord.
 
 ## quick start
 
-当前入门 llm.c 仓库的最佳方式是复现 GPT-2 (124M) 模型。[Discussion #481](https://github.com/karpathy/llm.c/discussions/481) 详细介绍了这一过程。我们可以在 llm.c 及并行的 PyTorch 实现中复现 GPT-2 和 GPT-3 系列的其他模型。请查看 [scripts README](scripts/README.md)。
+The best introduction to the llm.c repo today is reproducing the GPT-2 (124M) model. [Discussion #481](https://github.com/karpathy/llm.c/discussions/481) steps through this in detail. We can reproduce other models from the GPT-2 and GPT-3 series in both llm.c and in the parallel implementation of PyTorch. Have a look at the [scripts README](scripts/README.md).
 
-debugging tip: 运行 `make` 命令构建二进制文件时，将其中的 `-O3` 替换为 `-g`，即可在你喜爱的 IDE（例如 vscode）中单步调试代码。
+debugging tip: when you run the `make` command to build the binary, modify it by replacing `-O3` with `-g` so you can step through the code in your favorite IDE (e.g. vscode).
 
 ## quick start (1 GPU, fp32 only)
 
-如果你不会在多节点上训练、对混合精度不感兴趣，并且想学习 CUDA，fp32（legacy）文件可能适合你。这些文件是 llm.c 早期历史中的「checkpoint」版本，被定格在那一刻。它们更简单、更可移植，也可能更容易理解。单 GPU fp32 代码可按如下方式运行：
+If you won't be training on multiple nodes, aren't interested in mixed precision, and are interested in learning CUDA, the fp32 (legacy) files might be of interest to you. These are files that were "checkpointed" early in the history of llm.c and frozen in time. They are simpler, more portable, and possibly easier to understand. Run the 1 GPU, fp32 code like this:
 
 ```bash
 chmod u+x ./dev/download_starter_pack.sh
@@ -25,7 +19,7 @@ make train_gpt2fp32cu
 ./train_gpt2fp32cu
 ```
 
-`download_starter_pack.sh` 脚本是快速上手的简便方式，会下载一批有助于入门的 .bin 文件，其中包含：1) 以 fp32、bfloat16 保存的 GPT-2 124M 模型，2) 单元测试中使用的「debug state」（小批量数据及目标激活值和梯度），3) GPT-2 分词器，以及 3) 已分词化的 [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) 数据集。或者，不运行 .sh 脚本，也可按以下步骤手动重新创建这些产物：
+The download_starter_pack.sh script is a quick & easy way to get started and it downloads a bunch of .bin files that help get you off the ground. These contain: 1) the GPT-2 124M model saved in fp32, in bfloat16, 2) a "debug state" used in unit testing (a small batch of data, and target activations and gradients), 3) the GPT-2 tokenizer, and 3) the tokenized [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset. Alternatively, instead of running the .sh script, you can re-create these artifacts manually as follows:
 
 ```bash
 pip install -r requirements.txt
@@ -35,7 +29,7 @@ python train_gpt2.py
 
 ## quick start (CPU)
 
-「穷到连一块 GPU 都没有」专区。你仍然可以欣赏 llm.c 的训练过程！但走不了太远。与上述 fp32 版本类似，CPU 版本是 llm.c 历史上更早的一个 checkpoint，当时它还只是 C 语言的简单参考实现。例如，你可以不从头训练，而是将 GPT-2 small (124M) 微调为输出类莎士比亚文本，示例如下：
+The "I am so GPU poor that I don't even have one GPU" section. You can still enjoy seeing llm.c train! But you won't go too far. Just like the fp32 version above, the CPU version is an even earlier checkpoint in the history of llm.c, back when it was just a simple reference implementation in C. For example, instead of training from scratch, you can finetune a GPT-2 small (124M) to output Shakespeare-like text, as an example:
 
 ```bash
 chmod u+x ./dev/download_starter_pack.sh
@@ -44,9 +38,9 @@ make train_gpt2
 OMP_NUM_THREADS=8 ./train_gpt2
 ```
 
-如果你更希望不运行 starter pack 脚本，如上一节所述，可运行 `python dev/data/tinyshakespeare.py` 再运行 `python train_gpt2.py`，复现完全相同的 .bin 文件和产物。
+If you'd prefer to avoid running the starter pack script, then as mentioned in the previous section you can reproduce the exact same .bin files and artifacts by running `python dev/data/tinyshakespeare.py` and then `python train_gpt2.py`.
 
-上述命令（1）下载已分词化的 [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) 数据集并下载 GPT-2 (124M) 权重，（3）在 C 中从它们初始化，在 tineshakespeare 上用 AdamW 训练 40 步（batch size 4，context length 仅 64），评估验证集损失并采样部分文本。坦白说，除非你有一块强劲的 CPU（并能在启动命令中调高 OMP 线程数），否则在 CPU 上训练 LLM 走不了太远，但这可能是个不错的演示/参考。在我的 MacBook Pro（Apple Silicon M3 Max）上，输出如下：
+The above lines (1) download an already tokenized [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset and download the GPT-2 (124M) weights, (3) init from them in C and train for 40 steps on tineshakespeare with AdamW (using batch size 4, context length only 64), evaluate validation loss, and sample some text. Honestly, unless you have a beefy CPU (and can crank up the number of OMP threads in the launch command), you're not going to get that far on CPU training LLMs, but it might be a good demo/reference. The output looks like this on my MacBook Pro (Apple Silicon M3 Max):
 
 ```
 [GPT-2]
@@ -84,31 +78,31 @@ Allay
 
 ## datasets
 
-`/dev/data/(dataset).py` 内的数据文件负责下载、分词并将 token 保存为可从 C 轻松读取的 .bin 文件。例如运行：
+The data files inside `/dev/data/(dataset).py` are responsible for downloading, tokenizing and saving the tokens to .bin files, readable easily from C. So for example when you run:
 
 ```bash
 python dev/data/tinyshakespeare.py
 ```
 
-我们会下载并对 [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) 数据集进行分词。输出如下：
+We download and tokenize the [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset. The output of this looks like this:
 
 ```
 writing 32,768 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_val.bin
 writing 305,260 tokens to ./dev/data/tinyshakespeare/tiny_shakespeare_train.bin
 ```
 
-.bin 文件包含一个短头部（1024 字节），随后是以 uint16 存储的 token 流，表示 GPT-2 分词器的 token id。更多数据集可在 `/dev/data` 中找到。
+The .bin files contain a short header (1024 bytes) and then a stream of tokens in uint16, indicating the token ids with the GPT-2 tokenizer. More datasets are available in `/dev/data`.
 
 ## test
 
-我还附带了一个简单的单元测试，用于确保 C 代码与 PyTorch 代码一致。以 CPU 为例，编译并运行：
+I am also attaching a simple unit test for making sure our C code agrees with the PyTorch code. On the CPU as an example, compile and run with:
 
 ```bash
 make test_gpt2
 ./test_gpt2
 ```
 
-该测试会加载由 train_gpt2.py 写入的 `gpt2_124M_debug_state.bin` 文件，执行前向传播，将 logits 和 loss 与 PyTorch 参考实现对比，然后进行 10 次 Adam 训练迭代并确保 loss 与 PyTorch 一致。要测试 GPU 版本，运行：
+This now loads the `gpt2_124M_debug_state.bin` file that gets written by train_gpt2.py, runs a forward pass, compares the logits and loss with the PyTorch reference implementation, then it does 10 iterations of training with Adam and makes sure the losses match PyTorch. To test the GPU version we run:
 
 ```bash
 # fp32 test (cudnn not supported)
@@ -117,19 +111,19 @@ make test_gpt2cu PRECISION=FP32 && ./test_gpt2cu
 make test_gpt2cu USE_CUDNN=1 && ./test_gpt2cu
 ```
 
-这会同时测试 fp32 路径和混合精度路径。测试应通过并打印 `overall okay: 1`。
+This tests both the fp32 path and the mixed precision path. The test should pass and print `overall okay: 1`.
 
 ## tutorial
 
-我在此附上了一份很短的教程：[doc/layernorm/layernorm.md](doc/layernorm/layernorm.md)。这是一份简单、分步的指南，讲解如何实现 GPT-2 模型的单层——layernorm 层。这是理解 C 中如何实现各层的好起点。
+I attached a very small tutorial here, in [doc/layernorm/layernorm.md](doc/layernorm/layernorm.md). It's a simple, step-by-step guide to implementing a single layer of the GPT-2 model, the layernorm layer. This is a good starting point to understand how the layers are implemented in C.
 
-**flash attention**。截至 2024 年 5 月 1 日，我们使用 cuDNN 的 Flash Attention。由于 cuDNN 会将编译时间从几秒延长到约一分钟，且该代码路径目前非常新，因此默认禁用。可按如下方式编译以启用：
+**flash attention**. As of May 1, 2024 we use the Flash Attention from cuDNN. Because cuDNN bloats the compile time from a few seconds to ~minute and this code path is right now very new, this is disabled by default. You can enable it by compiling like this:
 
 ```bash
 make train_gpt2cu USE_CUDNN=1
 ```
 
-这将尝试使用 cudnn 编译并运行。你的系统上必须已安装 cuDNN。[cuDNN installation instructions](https://developer.nvidia.com/cudnn) 中通过 apt-get 安装会获取默认的 cuDNN 包集。最小化安装时，cuDNN 开发包即可，例如在 Ubuntu 22.04、CUDA 12.x 上：
+This will try to compile with cudnn and run it. You have to have cuDNN installed on your system. The [cuDNN installation instructions](https://developer.nvidia.com/cudnn) with apt-get will grab the default set of cuDNN packages. For a minimal setup, the cuDNN dev package is sufficient, e.g. on Ubuntu 22.04 for CUDA 12.x:
 
 ```bash
 wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
@@ -138,45 +132,45 @@ sudo apt-get update
 sudo apt-get -y install libcudnn9-dev-cuda-12
 ```
 
-此外还需要 [cuDNN frontend](https://github.com/NVIDIA/cudnn-frontend/tree/main),，但这只是头文件。将仓库克隆到本地即可。Makefile 目前会在你的主目录或当前目录中查找它。若放在其他位置，请在 `make` 命令行中添加 `CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include`。
+On top of this you need the [cuDNN frontend](https://github.com/NVIDIA/cudnn-frontend/tree/main), but this is just header files. Simply clone the repo to your disk. The Makefile currently looks for it in either your home directory or the current directory. If you have put it elsewhere, add `CUDNN_FRONTEND_PATH=/path/to/your/cudnn-frontend/include` to the `make` command-line.
 
-## 多 GPU 训练
+## multi-GPU training
 
-请确保已安装 MPI 和 NCCL，例如在 Linux 上：
+Make sure you install MPI and NCCL, e.g. on Linux:
 
 ```bash
 sudo apt install openmpi-bin openmpi-doc libopenmpi-dev
 ```
 
-对于 NCCL，请按照[官方网站](https://developer.nvidia.com/nccl/nccl-download)（例如网络安装器）中的说明操作
+For NCCL follow the instructions from the [official website](https://developer.nvidia.com/nccl/nccl-download) (e.g. network installer)
 
-然后：
+and then:
 
 ```bash
 make train_gpt2cu
 mpirun -np <number of GPUs> ./train_gpt2cu
 ```
 
-或者直接运行 `./scripts/` 下的某个脚本。
+or simply run one of our scripts under `./scripts/`.
 
-## 多节点训练
+## multi-node training
 
-请确保已按照 [multi-GPU](#multi-gpu-training) 章节的说明安装 `NCCL`。
+Make sure you've installed `NCCL` following instructions from [multi-GPU](#multi-gpu-training) section.
 
-目前我们支持 3 种方式来进行多节点训练：
-1) 使用 OpenMPI 交换 nccl id 并初始化 NCCL。详见例如 `./scripts/multi_node/run_gpt2_124M_mpi.sh` 脚本。
-2) 使用共享文件系统初始化 NCCL。详见 `./scripts/multi_node/run_gpt2_124M_fs.sbatch` 脚本。
-3) 使用 TCP 套接字初始化 NCCL。详见 `./scripts/multi_node/run_gpt2_124M_tcp.sbatch` 脚本。
+There are 3 ways we currently support that allow you to run multi-node training:
+1) Use OpenMPI to exchange nccl id and initialize NCCL. See e.g. `./scripts/multi_node/run_gpt2_124M_mpi.sh` script for details.
+2) Use shared file system to init NCCL. See `./scripts/multi_node/run_gpt2_124M_fs.sbatch` script for details.
+3) Use TCP sockets to init NCCL. See `./scripts/multi_node/run_gpt2_124M_tcp.sbatch` script for details.
 
-注意：
-* 如果你在 slurm 环境中运行，且你的 slurm 不支持 PMIx（考虑到 `slurm-wlm` 已移除 PMIx 支持，我们认为这会是常见情况），则必须使用 FS（2）或 TCP（3）方式。要测试你的 slurm 是否支持 PMIx，请运行：`srun --mpi=list`，并查看输出中是否出现 `pmix`。
-* 如果你没有配置 slurm，可以使用 `mpirun` - MPI（1）方式启动多节点运行。
+Note:
+* If you're running in a slurm environment and your slurm doesn't support PMIx (which we assume will be a common situation given that `slurm-wlm` dropped PMIx support) you will have to use FS (2) or TCP (3) approach. To test whether your slurm supports PMIx run: `srun --mpi=list` and see whether you get `pmix` in the output.
+* If you don't have slurm set up, you can kick off a multi-node run using `mpirun` - MPI (1).
 
-这 3 种方法并无优劣之分，我们只是提供多种选项，以便你在特定环境中运行。
+None of these 3 methods is superior, we just offer you options so that you can run in your specific environment.
 
-## 实验 / 参数扫描（sweeps）
+## experiments / sweeps
 
-下面是一个示例流程：在一台配备 4 个 GPU 的机器上，对 TinyStories 进行学习率扫描。运行 shell 脚本 `sweep.sh`（当然，在此之前你需要先 `chmod u+x sweep.sh`）：
+Just as an example process to sweep learning rates on a machine with 4 GPUs on TinyStories. Run a shell script `sweep.sh` (after you of course `chmod u+x sweep.sh`):
 
 ```bash
 #!/bin/bash
@@ -192,82 +186,64 @@ done
 # screen -ls | grep -E "tr[0-3]" | cut -d. -f1 | xargs -I {} screen -X -S {} quit
 ```
 
-此示例会开启 4 个 screen 会话，并以不同学习率运行四条命令。这会写出日志文件 `stories$i.log`，其中包含所有 loss，你可以用 Python 按需绘制。解析并绘制这些日志文件的快速示例见 [dev/vislog.ipynb](dev/vislog.ipynb)。
+This example opens up 4 screen sessions and runs the four commands with different LRs. This writes the log files `stories$i.log` with all the losses, which you can plot as you wish in Python. A quick example of how to parse and plot these logfiles is in [dev/vislog.ipynb](dev/vislog.ipynb).
 
-## 仓库
+## repo
 
-关于我希望这个仓库成为什么样子，再多说几句：
+A few more words on what I want this repo to be:
 
-首先，我希望 `llm.c` 成为一个教学场所。例如，我们的 `dev/cuda` 文件夹是一个内核库，涵盖各层的手工编写、文档详尽的内核，从非常简单的内核一直到更复杂/更快的内核。如果你有具有不同权衡（tradeoffs）的新内核，欢迎在此贡献。
+First, I want `llm.c` to be a place for education. E.g. our `dev/cuda` folder is a place for a library of kernels for all the layers that are manually hand-written and very well documented, starting from very simple kernels all the way to more complex / faster kernels. If you have a new kernel with various different tradeoffs, please feel free to contribute it here.
 
-话虽如此，我也希望 `llm.c` 非常快，甚至能实际用于训练网络。例如，起步时我们应该能够复现大型 GPT-2（1.6B）训练运行。这要求纳入最快的内核，包括使用 cuBLAS、cuBLASLt、CUTLASS、cuDNN 等库。我认为这样做也有教学意义：建立一个专家级上限和度量单位，例如你可以说你的手写内核达到 cuBLAS 速度的 80% 等。然后你可以选择超快运行，也可以选择“拖放”你想使用的任意手工内核来运行。
+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.
 
-不过，作为约束，我希望保持根目录中的主线 `llm.c` 简单且可读。如果某个 PR 例如将性能提升 2%，但“代价”是 500 行复杂 C 代码，也许还有冷门第三方依赖，我可能会拒绝该 PR，因为复杂度不值得。举个具体例子——在根训练循环中将 cuBLAS 用于 matmul 设为默认是显而易见的：它让主线代码快得多，只是一行可理解的代码，且是非常常见的依赖。与此同时，我们可以在 `dev/cuda` 中提供能与 cuBLAS 竞争的手工实现。
+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`.
 
-最后，我对项目根文件夹中的复杂度会敏感得多，根目录包含项目的主/默认文件。相比之下，`dev/` 文件夹更像是我们开发内核或类库、分享有用/相关/教学代码的草稿空间，其中部分代码可以（在局部）复杂一些。
+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.
 
-## 值得关注的分支（fork）
+## notable forks
 
-- AMD 支持
-  - [llm.c](https://github.com/anthonix/llm.c) by @[anthonix](https://github.com/anthonix): 支持 AMD 设备，例如 7900 XTX
+- AMD support
+  - [llm.c](https://github.com/anthonix/llm.c) by @[anthonix](https://github.com/anthonix): support for AMD devices, such as the 7900 XTX
 
 - C#
-  - [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): 本项目的 C# 移植版
-  - [Llm.cs](https://github.com/nietras/Llm.cs) by @[nietras](https://github.com/nietras): 本项目的 C# 移植版，侧重于在任何平台上易于上手。克隆即可运行 ✅
+  - [llm.cs](https://github.com/azret/llm.cs) by @[azret](https://github.com/azret): a C# port of this project
+  - [Llm.cs](https://github.com/nietras/Llm.cs) by @[nietras](https://github.com/nietras): a C# port of this project with focus on easy to get started on any platform. Clone and run ✅
 
 - CUDA C++
-  - [llm.cpp](https://github.com/gevtushenko/llm.c) by @[gevtushenko](https://github.com/gevtushenko): 使用 [CUDA C++ Core Libraries](https://github.com/NVIDIA/cccl) 对本项目的移植
-     - 该分支曾在 [GPU MODE Discord Server](https://discord.gg/cudamode) 的[这场讲座](https://www.youtube.com/watch?v=WiB_3Csfj_Q) 中介绍
-
-- C++/CUDA
-  - [llm.cpp](https://github.com/zhangpiu/llm.cpp/tree/master/llmcpp) by @[zhangpiu](https://github.com/zhangpiu): 使用 [Eigen](https://gitlab.com/libeigen/eigen), 对本项目的移植，支持 CPU/CUDA。
-
-- WebGPU C++
-  - [gpu.cpp](https://github.com/AnswerDotAI/gpu.cpp) by @[austinvhuang](https://github.com/austinvhuang): 使用原生 WebGPU 在 C++ 中进行可移植 GPU 计算的库。目标是成为通用库，同时也将 llm.c 内核移植到 WGSL。
-  
-- C++
-  - [llm.cpp](https://github.com/GaoYusong/llm.cpp) by @[GaoYusong](https://github.com/GaoYusong): 本项目的 C++ 移植版，提供 C++ 单头文件库 [tinytorch.hpp](https://github.com/GaoYusong/llm.cpp/blob/main/tinytorch.hpp)
+  - [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)
 
 - Go
-  - [llm.go](https://github.com/joshcarp/llm.go) by @[joshcarp](https://github.com/joshcarp): 本项目的 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): 本项目的 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): 使用简洁、原生 C/Metal Shading Language 进行 LLM 训练
+  - [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): 本项目的 Mojo 移植版
-
-- OpenCL
-  - [llm.c](https://github.com/krrishnarraj/llm.c) by @[krrishnarraj](https://github.com/krrishnarraj): 本项目的 OpenCL 移植版
+  - [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/yijunyu/llm.rs) by @[Yijun Yu](https://github.com/yijunyu): Rust 重写版，目标是达到相同性能
-  -  [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): 本项目的 Rust 移植版
+  -  [llm.rs](https://github.com/yijunyu/llm.rs) by @[Yijun Yu](https://github.com/yijunyu): a Rust rewrite with the aim to have same performance
+  -  [llm.rs](https://github.com/ToJen/llm.rs) by @[ToJen](https://github.com/ToJen): a Rust port of this project
 
 - Swift
-  - [llm.swift](https://github.com/otabuzzman/llm.swift) by @[otabuzzman](https://github.com/otabuzzman): 本项目的 Swift 移植版
+  - [llm.swift](https://github.com/otabuzzman/llm.swift) by @[otabuzzman](https://github.com/otabuzzman): a Swift port of this project
 
 - Zig
-  - [llm.zig](https://github.com/Saimirbaci/llm.zig) by @[saimirbaci](https://github.com/Saimirbaci): 本项目的 Zig 移植版
- 
-- Habana Gaudi2
-  - [llm.tpc](https://github.com/abhilash1910/llm.tpc) by @[abhilash1910](https://github.com/abhilash1910): 本项目的 Habana Gaudi2 移植版 
+  - [llm.zig](https://github.com/Saimirbaci/llm.zig) by @[saimirbaci](https://github.com/Saimirbaci): a Zig port of this project
 
-- Nim
-  - [llm.nim](https://github.com/Vindaar/llm.nim) by @[Vindaar](https://github.com/Vindaar): 本项目的 Nim 移植版
+## discussions
 
-## 讨论
+Ways of organizing development:
 
-组织开发的方式：
+- Experiencing a concrete issue with the repo? Use [Issues](https://github.com/karpathy/llm.c/issues).
+- Have some code to contribute? Open a [PR](https://github.com/karpathy/llm.c/pulls)
+- Chat about the repo, ask questions, etc.? Look at [Discussions](https://github.com/karpathy/llm.c/discussions).
+- Something faster? I created a new `#llmc` channel on my [Zero to Hero Discord channel](https://discord.gg/3zy8kqD9Cp).
 
-- 在使用仓库时遇到具体问题？请使用 [Issues](https://github.com/karpathy/llm.c/issues).
-- 有代码要贡献？请提交 [PR](https://github.com/karpathy/llm.c/pulls)
-- 想聊聊仓库、提问等？请查看 [Discussions](https://github.com/karpathy/llm.c/discussions).
-- 需要更快捷的沟通？我在我的 [Zero to Hero Discord 频道](https://discord.gg/3zy8kqD9Cp). 上新建了 `#llmc` 频道
-
-## 许可证
+## license
 
 MIT

README.wehub.md

@@ -1,7 +0,0 @@
-# WeHub 来源说明
-
-- 原始项目：`karpathy/llm.c`
-- 原始仓库：https://github.com/karpathy/llm.c
-- 导入方式：上游默认分支的最新快照
-- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
-- 本文件仅用于记录来源，不代表 WeHub 是原项目作者

dev/cuda/Makefile

@@ -8,20 +8,8 @@ ifeq ($(NVCC),)
 		$(error nvcc not found.)
 endif
 
-ifneq ($(CI),true) # if not in CI, then use the GPU query
-  ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY=
-    GPU_COMPUTE_CAPABILITY = $(shell __nvcc_device_query) # assume if NVCC is present, then this likely is too
-    GPU_COMPUTE_CAPABILITY := $(strip $(GPU_COMPUTE_CAPABILITY))
-  endif
-endif
-
 # Compiler flags
-ifeq ($(GPU_COMPUTE_CAPABILITY),) # set to defaults if: make GPU_COMPUTE_CAPABILITY=
-  CFLAGS = -O3 --use_fast_math
-else
-  CFLAGS = -O3 --use_fast_math --generate-code arch=compute_$(GPU_COMPUTE_CAPABILITY),code=[compute_$(GPU_COMPUTE_CAPABILITY),sm_$(GPU_COMPUTE_CAPABILITY)]
-endif
-
+CFLAGS = -O3 --use_fast_math
 NVCCFLAGS = -lcublas -lcublasLt -std=c++17
 MPI_PATHS = -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux-gnu/openmpi/lib/
 
@@ -30,11 +18,8 @@ MPI_PATHS = -I/usr/lib/x86_64-linux-gnu/openmpi/include -L/usr/lib/x86_64-linux-
 	$(NVCC) $(CFLAGS) $(NVCCFLAGS) $< -o $@
 
 # Build all targets
-TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward fused_residual_forward  global_norm permute
-
+TARGETS = adamw attention_backward attention_forward classifier_fused crossentropy_forward crossentropy_softmax_backward encoder_backward encoder_forward gelu_backward gelu_forward layernorm_backward layernorm_forward matmul_backward matmul_backward_bias matmul_forward nccl_all_reduce residual_forward softmax_forward trimat_forward fused_residual_forward  global_norm
 all: $(TARGETS)
-all_ptx:  $(TARGETS:%=%.ptx)
-all_sass: $(TARGETS:%=%.sass)
 
 # Individual targets: forward pass
 attention_forward: attention_forward.cu
@@ -65,20 +50,10 @@ matmul_backward: matmul_backward.cu
 adamw: adamw.cu
 global_norm: global_norm.cu
 
-permute: permute.cu
-
 # NCCL communication kernels
 nccl_all_reduce: nccl_all_reduce.cu
 	$(NVCC) -lmpi -lnccl $(NVCCFLAGS) $(MPI_PATHS) nccl_all_reduce.cu -o nccl_all_reduce
 
-# Generate PTX using cuobjdump
-%.ptx: %
-	cuobjdump --dump-ptx $< > $@
-
-# Generate SASS using cuobjdump
-%.sass: %
-	cuobjdump --dump-sass $< > $@
-
 # Run all targets
 run_all: all
 	@for target in $(TARGETS); do \
@@ -90,4 +65,4 @@ run_all: all
 
 # Clean up
 clean:
-	rm -f $(TARGETS) *.ptx *.sass
+	rm -f $(TARGETS)

dev/cuda/attention_backward.cu

@@ -1137,7 +1137,6 @@ int main(int argc, char **argv) {
     free(dinp);
     free(dpreatt);
     free(datt);
-    free(h_dinp);
     cudaCheck(cudaFree(d_inp));
     cudaCheck(cudaFree(d_qkvr));
     cudaCheck(cudaFree(d_preatt));

dev/cuda/attention_forward.cu

@@ -1377,7 +1377,6 @@ int main(int argc, char **argv) {
     cudaCheck(cudaFree(d_preatt));
     cudaCheck(cudaFree(d_att));
     cudaCheck(cudaFree(d_inp));
-    cudaCheck(cudaFree(d_stats));
     cublasDestroy(cublas_handle);
 
     #ifdef ENABLE_CUDNN

dev/cuda/classifier_fused.cu

@@ -766,7 +766,6 @@ int main(int argc, char **argv) {
     cudaCheck(cudaFree(d_logits));
     cudaCheck(cudaFree(d_dlosses));
     cudaCheck(cudaFree(d_targets));
-    cudaCheck(cudaFree(d_dlogits_no_pad));
 
     return 0;
 }

dev/cuda/layernorm_forward.cu

@@ -28,6 +28,9 @@ verstion 5 allocates blocks per row instead of warps per row, same alg as 4 othe
 #include <cooperative_groups.h>
 #include <cooperative_groups/reduce.h>
 #include "common.h"
+
+#define WARP_SIZE 32
+
 // ----------------------------------------------------------------------------
 // CPU code reference
 

dev/cuda/nccl_all_reduce.cu

@@ -193,6 +193,5 @@ int main(int argc, char **argv) {
 
   free(all_reduce_buffer_host);
   cudaCheck(cudaFree(all_reduce_buffer));
-  cudaCheck(cudaFree(all_reduce_buffer_recv));
   multi_gpu_config_free(&multi_gpu_config);
 }

dev/cuda/permute.cu

@@ -1,181 +0,0 @@
-/*
-Kernels to demonstrate permute operation.
-
-Compile example:
-nvcc -O3 permute.cu -o permute
-
-The goal is to permute a 4D matrix from its original shape (dim1, dim2, dim3, dim4) to a new shape (dim4, dim3, dim1, dim2).
-
-Before permutation, we need to understand how to access elements in a flattened (linear) form of the matrix.
-
-Given:
-
-dim1 = size of the 1st dimension
-dim2 = size of the 2nd dimension
-dim3 = size of the 3rd dimension
-dim4 = size of the 4th dimension
-
-For any element in a 4D matrix at position (i1, i2, i3, i4), where:
-
-i1 is the index in dimension 1
-i2 is the index in dimension 2
-i3 is the index in dimension 3
-i4 is the index in dimension 4
-
-If you find it challenging to calculate the indices i1, i2, i3, and i4, observe the pattern in the index calculations.
-Initially, it might take some time to grasp, but with practice, you'll develop a mental model for it.
-
-To calculate the indices, use the following formulas:
-
-i1 = (idx / (dim2 * dim3 * dim4)) % dim1;
-i2 = (idx / (dim3 * dim4)) % dim2;
-i3 = (idx / dim4) % dim3;
-i4 = idx % dim4;
-
-Pattern Explanation:
-To find the index for any dimension, divide the thread ID (idx) by the product of all subsequent dimensions.
-Then, perform modulo operation with the current dimension.
-
-
-
-The linear index in a flattened 1D array is calculated as:
-linear_idx = i1 × ( dim2 × dim3 × dim4 ) + i2 × ( dim3 × dim4 ) + i3 × dim4 + i4
-This linear index uniquely identifies the position of the element in the 1D array.
-
-To permute the matrix, we need to rearrange the indices according to the new shape.
-In this case, we are permuting from (dim1, dim2, dim3, dim4) to (dim4, dim3, dim1, dim2).
-
-The new dimension post permutation will be as follows:
-
-dim1 becomes the new 3rd dimension.
-dim2 becomes the new 4th dimension.
-dim3 becomes the new 2nd dimension.
-dim4 becomes the new 1st dimension.
-
-permuted_idx = i4 * (dim3 * dim1 * dim2) + i3 * (dim1 * dim2) + i1 * dim2 + i2;
-
-Here's how this works:
-
-i4 * (dim3 * dim1 * dim2): This accounts for how many complete dim3 × dim1 × dim2 blocks fit before the current i4 block.
-i3 * (dim1 * dim2): This accounts for the offset within the current i4 block, specifying which i3 block we are in.
-i1 * dim2: This accounts for the offset within the current i3 block, specifying which i1 block we are in.
-i2: This gives the offset within the current i1 block.
-
-Lastly at the end we store the current value at idx index of the original value to the permuted index in the permuted_matrix.
-
-
---------------------------------------------------------------------------------------------------------------------------------------------------------
-
-Similarly we can follow the above approach to permute matrices of any dimensions.
-
-*/
-
-
-#include <cuda_runtime.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <cmath>
-
-#include "common.h"
-
-// CPU function to permute a 4D matrix
-void permute_cpu(const float* matrix, float* out_matrix, int dim1, int dim2, int dim3, int dim4) {
-    int total_threads = dim1 * dim2 * dim3 * dim4;
-
-    for (int idx = 0; idx < total_threads; idx++) {
-        // Calculate the 4D indices from the linear index
-        int i1 = (idx / (dim2 * dim3 * dim4)) % dim1;
-        int i2 = (idx / (dim3 * dim4)) % dim2;
-        int i3 = (idx / dim4) % dim3;
-        int i4 = idx % dim4;
-
-        // Compute the new index for the permuted matrix
-        // Transpose from (dim1, dim2, dim3, dim4) to (dim4, dim3, dim1, dim2)
-        int permuted_idx = i4 * (dim3 * dim1 * dim2) + i3 * (dim1 * dim2) + i1 * dim2 + i2;
-        out_matrix[permuted_idx] = matrix[idx];
-    }
-}
-
-// CUDA kernel to permute a 4D matrix
-__global__ void permute_kernel(const float* matrix, float* out_matrix, int dim1, int dim2, int dim3, int dim4) {
-    int idx = blockIdx.x * blockDim.x + threadIdx.x;
-
-    // Ensure index is within bounds
-    if (idx < dim1 * dim2 * dim3 * dim4) {
-        // Calculate the 4D indices from the linear index
-        int i1 = (idx / (dim2 * dim3 * dim4)) % dim1;
-        int i2 = (idx / (dim3 * dim4)) % dim2;
-        int i3 = (idx / dim4) % dim3;
-        int i4 = idx % dim4;
-
-        // Compute the new index for the permuted matrix
-        // Transpose from (dim1, dim2, dim3, dim4) to (dim4, dim3, dim1, dim2)
-        int permuted_idx = i4 * (dim3 * dim1 * dim2) + i3 * (dim1 * dim2) + i1 * dim2 + i2;
-        out_matrix[permuted_idx] = matrix[idx];
-    }
-}
-
-
-int main() {
-    int dim_1 = 24;
-    int dim_2 = 42;
-    int dim_3 = 20;
-    int dim_4 = 32;
-
-    // Set up the device
-    int deviceIdx = 0;
-    cudaSetDevice(deviceIdx);
-    cudaDeviceProp deviceProp;
-    cudaGetDeviceProperties(&deviceProp, deviceIdx);
-    printf("Device %d: %s\n", deviceIdx, deviceProp.name);
-
-    // Allocate host memory
-    float* matrix = make_random_float(dim_1 * dim_2 * dim_3 * dim_4);
-    float* permuted_matrix = (float*)malloc(dim_1 * dim_2 * dim_3 * dim_4 * sizeof(float));
-
-    // Initialize the matrix with random values
-
-    // Allocate device memory
-    float *d_matrix, *d_permuted_matrix;
-    cudaMalloc(&d_matrix, dim_1 * dim_2 * dim_3 * dim_4 * sizeof(float));
-    cudaMalloc(&d_permuted_matrix, dim_1 * dim_2 * dim_3 * dim_4 * sizeof(float));
-
-    // Copy matrix from host to device
-    cudaMemcpy(d_matrix, matrix, dim_1 * dim_2 * dim_3 * dim_4 * sizeof(float), cudaMemcpyHostToDevice);
-
-    // Perform permutation on CPU
-    clock_t start = clock();
-    permute_cpu(matrix, permuted_matrix, dim_1, dim_2, dim_3, dim_4);
-    clock_t end = clock();
-    double elapsed_time_cpu = (double)(end - start) / CLOCKS_PER_SEC;
-
-    // Define block and grid sizes
-    dim3 blockSize(256);
-    int totalThreads = dim_1 * dim_2 * dim_3 * dim_4;
-    int gridSize = (totalThreads + blockSize.x - 1) / blockSize.x; // Compute grid size
-
-    // Launch CUDA kernel to perform permutation
-    permute_kernel<<<gridSize, blockSize>>>(d_matrix, d_permuted_matrix, dim_1, dim_2, dim_3, dim_4);
-    cudaDeviceSynchronize(); // Ensure kernel execution is complete
-
-    // Verify results
-    printf("Checking correctness...\n");
-    validate_result(d_permuted_matrix, permuted_matrix, "permuted_matrix", dim_1 * dim_2 * dim_3 * dim_4, 1e-5f);
-
-    printf("All results match.\n\n");
-    // benchmark kernel
-    int repeat_times = 1000;
-    float elapsed_time = benchmark_kernel(repeat_times, permute_kernel,
-                                          d_matrix, d_permuted_matrix, dim_1, dim_2, dim_3, dim_4
-    );
-    printf("time gpu %.4f ms\n", elapsed_time);
-    printf("time cpu %.4f ms\n", elapsed_time_cpu);
-
-    // Free allocated memory
-    free(matrix);
-    free(permuted_matrix);
-    cudaFree(d_matrix);
-    cudaFree(d_permuted_matrix);
-
-    return 0;
-}

dev/cuda/trimat_forward.cu

@@ -643,7 +643,6 @@ int main(int argc, char **argv) {
     free(inp);
     cudaCheck(cudaFree(d_out));
     cudaCheck(cudaFree(d_inp));
-    cudaCheck(cudaFree(d_qkvr));
     cublasDestroy(cublas_handle);
 
     return 0;

dev/data/README.md

@@ -6,5 +6,3 @@ The idea is that each dataset has a .py file here in the root of `dev/data`, and
 - running `python tinyshakespeare.py` will create a directory `tinyshakespeare` with its .bin files inside it
 
 And so on. This way we can nicely organize multiple datasets here, share common utilities between them, and then point the .py/.c code in the root of the project accordingly to these.
-
-Note: we support "gpt-2" and "llama" (llama 3 in particular) models and the above scripts will tokenize gpt-2 by default.

dev/data/data_common.py

@@ -23,38 +23,28 @@ def download_file(url: str, fname: str, chunk_size=1024):
             bar.update(size)
 
 
-HEADERS_INFO = {
-    "gpt-2": {
-        "magic": 20240520,
-        "version": 1,
-        "token_dtype": np.uint16,
-    },
-    "llama-3": {
-        "magic": 20240801,
-        "version": 7,
-        "token_dtype": np.uint32,
-    },
-}
-
-def write_datafile(filename, toks, model_desc="gpt-2"):
+def write_datafile(filename, toks):
     """
     Saves token data as a .bin file, for reading in C.
     - First comes a header with 256 int32s
-    - The tokens follow, each as uint16 (gpt-2) or uint32 (llama)
+    - The tokens follow, each as a uint16
     """
     assert len(toks) < 2**31, "token count too large" # ~2.1B tokens
-    assert model_desc in ["gpt-2", "llama-3"], f"unknown model descriptor {model_desc}"
-    info = HEADERS_INFO[model_desc]
     # construct the header
-    header = np.zeros(256, dtype=np.int32) # header is always 256 int32 values
-    header[0] = info["magic"]
-    header[1] = info["version"]
-    header[2] = len(toks) # number of tokens after the 256*4 bytes of header
-    # construct the data (numpy array of tokens)
-    toks_np = np.array(toks, dtype=info["token_dtype"])
+    header = np.zeros(256, dtype=np.int32)
+    header[0] = 20240520 # magic
+    header[1] = 1 # version
+    header[2] = len(toks) # number of tokens after the 256*4 bytes of header (each 2 bytes as uint16)
+    # construct the tokens numpy array, if not already
+    if not isinstance(toks, np.ndarray) or not toks.dtype == np.uint16:
+        # validate that no token exceeds a uint16
+        maxtok = 2**16
+        assert all(0 <= t < maxtok for t in toks), "token dictionary too large for uint16"
+        toks_np = np.array(toks, dtype=np.uint16)
+    else:
+        toks_np = toks
     # write to file
-    num_bytes = (256 * 4) + (len(toks) * toks_np.itemsize)
-    print(f"writing {len(toks):,} tokens to {filename} ({num_bytes:,} bytes) in the {model_desc} format")
+    print(f"writing {len(toks):,} tokens to {filename}")
     with open(filename, "wb") as f:
         f.write(header.tobytes())
         f.write(toks_np.tobytes())

dev/data/edu_fineweb.sh

@@ -1,73 +0,0 @@
-#!/bin/bash
-
-# Downloads the FineWeb-Edu 100B dataset, but in an already tokenized format in .bin files
-# Example: ./edu_fineweb.sh 100
-# would download 100 shards
-# Default is all shards
-# Make sure to run this from current directory, i.e. inside ./dev/data!
-
-# Check if MAX_SHARDS is provided as positional first arg, otherwise default to 1024
-if [ $# -eq 0 ]; then
-    MAX_SHARDS=1001
-else
-    MAX_SHARDS=$1
-fi
-
-if [ $MAX_SHARDS -gt 1001 ]; then
-    MAX_SHARDS=1001
-fi
-
-# Base URLs
-TRAIN_BASE_URL="https://huggingface.co/datasets/karpathy/fineweb-edu-100B-gpt2-token-shards/resolve/main/edu_fineweb_train_"
-VAL_URL="https://huggingface.co/datasets/karpathy/fineweb-edu-100B-gpt2-token-shards/resolve/main/edu_fineweb_val_000000.bin"
-
-# Directory to save files
-SAVE_DIR="edu_fineweb100B"
-
-# Create the directory if it doesn't exist
-mkdir -p "$SAVE_DIR"
-
-download() {
-    local FILE_URL=$1
-    local FILE_NAME=$(basename $FILE_URL | cut -d'?' -f1)
-    local FILE_PATH="${SAVE_DIR}/${FILE_NAME}"
-    curl -s -L -o "$FILE_PATH" "$FILE_URL"
-    echo "Downloaded $FILE_NAME to $SAVE_DIR"
-}
-
-# Function to manage parallel jobs
-run_in_parallel() {
-    local max_jobs=$1
-    shift
-    local commands=("$@")
-    local job_count=0
-
-    for cmd in "${commands[@]}"; do
-        eval "$cmd" &
-        ((job_count++))
-        if (( job_count >= max_jobs )); then
-            wait -n
-            ((job_count--))
-        fi
-    done
-
-    # Wait for any remaining jobs to finish
-    wait
-}
-
-# Export the function so it's available in subshells
-export -f download
-
-# Download the validation shard
-download "$VAL_URL" &
-
-# Generate train file shard download commands
-train_commands=()
-for i in $(seq -f "%06g" 1 $MAX_SHARDS); do
-    FILE_URL="${TRAIN_BASE_URL}${i}.bin?download=true"
-    train_commands+=("download \"$FILE_URL\"")
-done
-
-# Run the train file commands in parallel
-run_in_parallel 40 "${train_commands[@]}"
-echo "The val shard and first $MAX_SHARDS train shards of FineWebEdu100B files downloaded in $SAVE_DIR"

dev/data/fineweb.py

@@ -14,102 +14,61 @@ example doc to highlight the structure of the dataset:
   "language_score": 0.9185474514961243,
   "token_count": 594
 }
-
-Example of downloading the 100B dataset of FineWebEDU, from root directory:
-python dev/data/fineweb.py -t edu -v 100B
-100B runs for small few hours, depending on your internet and computer.
 """
 import os
 import argparse
 import multiprocessing as mp
-
 import numpy as np
 import tiktoken
+# from huggingface_hub import snapshot_download
 from datasets import load_dataset
 from tqdm import tqdm
-
-from transformers import AutoTokenizer
-
+import argparse
 
 from data_common import write_datafile
 # ------------------------------------------
 
-parser = argparse.ArgumentParser(description="FineWeb and Edu-FineWeb dataset preprocessing")
-parser.add_argument("-t", "--type", type=str, default="classic", help="Fineweb type, edu|classic")
-parser.add_argument("-v", "--version", type=str, default="10B", help="Fineweb data sample size, 10B|100B")
-parser.add_argument("-m", "--model_desc", type=str, default="gpt-2", help="Model descriptor, gpt-2|llama-3")
-parser.add_argument("-s", "--shard_size", type=int, default=10**8, help="Size of each data shard in the output .bin files, in tokens")
+parser = argparse.ArgumentParser(description="FineWeb dataset preprocessing")
+parser.add_argument("-v", "--version", type=str, default="10B", help="Which version of fineweb to use 10B|100B")
+parser.add_argument("-s", "--shard_size", type=int, default=10**8, help="Size of each shard in tokens")
 args = parser.parse_args()
 
 # FineWeb has a few possible subsamples available
-assert args.version in {"10B", "100B"}, "version must be one of: 10B, 100B"
-assert args.type in {"edu", "classic"}, "type must be one of: edu, classic"
-directories = {
-    ("classic", "10B"): ("fineweb10B", "sample-10BT"),
-    ("classic", "100B"): ("fineweb100B", "sample-100BT"),
-    ("edu", "10B"): ("edu_fineweb10B", "sample-10BT"),
-    ("edu", "100B"): ("edu_fineweb100B", "sample-100BT")
-}
-local_dir, remote_name = directories[(args.type, args.version)]
+assert args.version in ["10B", "100B"], "version must be one of 10B, 100B"
+if args.version == "10B":
+    local_dir = "fineweb10B"
+    remote_name = "sample-10BT"
+elif args.version == "100B":
+    local_dir = "fineweb100B"
+    remote_name = "sample-100BT"
 
 # create the cache the local directory if it doesn't exist yet
 DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), local_dir)
 os.makedirs(DATA_CACHE_DIR, exist_ok=True)
 
 # download the dataset
-if args.type == "classic":
-    fw = load_dataset("HuggingFaceFW/fineweb", name=remote_name, split="train")
-    name = "fineweb"
-elif args.type =="edu":
-    fw = load_dataset("HuggingFaceFW/fineweb-edu", name=remote_name, split="train")
-    name = "edu_fineweb"
+fw = load_dataset("HuggingFaceFW/fineweb", name=remote_name, split="train")
 
-def tokenize_llama(doc):
-    # tokenizes a single document and returns a numpy array of uint32 tokens
-    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
-    encode = lambda s: tokenizer.encode(s, add_special_tokens=False, verbose=False, split_special_tokens=True)
-    eot = tokenizer.encode('')[0] # by default the tokenizer adds the EOT token (128000)
-    tokens = [eot] # the special <|endoftext|> token delimits all documents
-    tokens.extend(encode(doc["text"]))
-    tokens_np = np.array(tokens)
-    assert (0 <= tokens_np).all() and (tokens_np < 2**32).all(), "token dictionary too large for uint32"
-    tokens_np_uint = tokens_np.astype(np.uint32)
-    return tokens_np_uint
-
-def tokenize_gpt2(doc):
+# init the tokenizer
+enc = tiktoken.get_encoding("gpt2")
+eot = enc._special_tokens['<|endoftext|>'] # end of text token
+def tokenize(doc):
     # tokenizes a single document and returns a numpy array of uint16 tokens
-    enc = tiktoken.get_encoding("gpt2")
-    encode = lambda s: enc.encode_ordinary(s)
-    eot = enc._special_tokens['<|endoftext|>'] # end of text token
     tokens = [eot] # the special <|endoftext|> token delimits all documents
-    tokens.extend(encode(doc["text"]))
+    tokens.extend(enc.encode_ordinary(doc["text"]))
     tokens_np = np.array(tokens)
     assert (0 <= tokens_np).all() and (tokens_np < 2**16).all(), "token dictionary too large for uint16"
-    tokens_np_uint = tokens_np.astype(np.uint16)
-    return tokens_np_uint
-
-token_dtype = {
-    "gpt-2": np.uint16,
-    "llama-3": np.uint32
-}[args.model_desc]
+    tokens_np_uint16 = tokens_np.astype(np.uint16)
+    return tokens_np_uint16
 
 # tokenize all documents and write output shards, each of shard_size tokens (last shard has remainder)
 nprocs = max(1, os.cpu_count() - 2) # don't hog the entire system
 with mp.Pool(nprocs) as pool:
     shard_index = 0
     # preallocate buffer to hold current shard
-    all_tokens_np = np.empty((args.shard_size,), dtype=token_dtype)
+    all_tokens_np = np.empty((args.shard_size,), dtype=np.uint16)
     token_count = 0
     progress_bar = None
-
-    tokenize = lambda x: None
-    if args.model_desc == "gpt-2":
-        tokenize = tokenize_gpt2
-    elif args.model_desc == "llama-3":
-        tokenize = tokenize_llama
-    else:
-        raise ValueError(f"unknown model {args.model_desc}")
-
     for tokens in pool.imap(tokenize, fw, chunksize=16):
 
         # is there enough space in the current shard for the new tokens?
@@ -124,12 +83,12 @@ with mp.Pool(nprocs) as pool:
         else:
             # write the current shard and start a new one
             split = "val" if shard_index == 0 else "train"
-            filename = os.path.join(DATA_CACHE_DIR, f"{name}_{split}_{shard_index:06d}.bin")
+            filename = os.path.join(DATA_CACHE_DIR, f"fineweb_{split}_{shard_index:06d}.bin")
             # split the document into whatever fits in this shard; the remainder goes to next one
             remainder = args.shard_size - token_count
             progress_bar.update(remainder)
             all_tokens_np[token_count:token_count+remainder] = tokens[:remainder]
-            write_datafile(filename, all_tokens_np.tolist(), args.model_desc)
+            write_datafile(filename, all_tokens_np)
             shard_index += 1
             progress_bar = None
             # populate the next shard with the leftovers of the current doc
@@ -139,5 +98,5 @@ with mp.Pool(nprocs) as pool:
     # write any remaining tokens as the last shard
     if token_count != 0:
         split = "val" if shard_index == 0 else "train"
-        filename = os.path.join(DATA_CACHE_DIR, f"{name}_{split}_{shard_index:06d}.bin")
-        write_datafile(filename, (all_tokens_np[:token_count]).tolist(), args.model_desc)
+        filename = os.path.join(DATA_CACHE_DIR, f"fineweb_{split}_{shard_index:06d}.bin")
+        write_datafile(filename, all_tokens_np[:token_count])

dev/data/tinyshakespeare.py

@@ -6,32 +6,25 @@ Downloads and tokenizes the TinyShakespeare dataset.
 The output is written to a newly created tinyshakespeare/ folder.
 The script prints:
 
-For GPT-2:
-$ python dev/data/tinyshakespeare.py --model=gpt-2
-writing 32,768 tokens to /home/ubuntu/llm.c/dev/data/tinyshakespeare/tiny_shakespeare_val.bin (66,560 bytes) in the gpt-2 format
-writing 305,260 tokens to /home/ubuntu/llm.c/dev/data/tinyshakespeare/tiny_shakespeare_train.bin (611,544 bytes) in the gpt-2 format
-
-For LLaMA 3:
-$ python dev/data/tinyshakespeare.py --model=llama-3
-writing 32,768 tokens to /home/ubuntu/llm.c/dev/data/tinyshakespeare/tiny_shakespeare_val.bin (132,096 bytes) in the llama-3 format
-writing 276,224 tokens to /home/ubuntu/llm.c/dev/data/tinyshakespeare/tiny_shakespeare_train.bin (1,105,920 bytes) in the llama-3 format
+Saved 32768 tokens to tinyshakespeare/tiny_shakespeare_val.bin
+Saved 305260 tokens to tinyshakespeare/tiny_shakespeare_train.bin
 
 And runs in a few seconds depending on your internet
 connection and computer. The .bin files are raw byte
-streams of uint16 (gpt-2) or uint32 (llama) numbers indicating the token ids.
+streams of int32 numbers indicating the token ids.
 """
 
-import argparse
 import os
-
 import tiktoken
-from transformers import AutoTokenizer
-
+import numpy as np
 from data_common import download_file, write_datafile
 
 # -----------------------------------------------------------------------------
 DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), "tinyshakespeare")
 
+enc = tiktoken.get_encoding("gpt2")
+encode = lambda s: enc.encode(s, allowed_special={'<|endoftext|>'})
+
 def download():
     """Downloads the TinyShakespeare dataset to DATA_CACHE_DIR"""
     os.makedirs(DATA_CACHE_DIR, exist_ok=True)
@@ -44,41 +37,23 @@ def download():
     else:
         print(f"{data_filename} already exists, skipping download...")
 
-def tokenize(model_desc):
-    if model_desc == "gpt-2":
-        enc = tiktoken.get_encoding("gpt2")
-        encode = lambda s: enc.encode_ordinary(s)
-        eot = enc._special_tokens['<|endoftext|>'] # end of text token
-    elif model_desc == "llama-3":
-        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
-        encode = lambda s: tokenizer.encode(s, add_special_tokens=False, verbose=False, split_special_tokens=True)
-        eot = tokenizer.encode('')[0] # by default the tokenizer adds the EOT token (128000)
-    else:
-        raise ValueError(f"unknown model descriptor {model_desc}")
+def tokenize():
     data_filename = os.path.join(DATA_CACHE_DIR, "tiny_shakespeare.txt")
     text = open(data_filename, 'r').read()
-    # let's treat every individual chunk of text as a separate "document"
-    sections = text.split("\n\n")
-    tokens = []
-    for i, s in enumerate(sections):
-        tokens.append(eot)
-        # there was a mild bug where I originally intended to remove \n\n, but instead just added
-        # the EOT right after each \n\n, so I'm keeping that behavior for backwards compatibility
-        # therefore we have to here add an extra \n\n at the end of each section, except the last
-        spad = s + "\n\n" if i != len(sections) - 1 else s
-        tokens.extend(encode(spad))
+    # let's treat every person's statement in the dialog as a separate document
+    text = "<|endoftext|>" + text
+    text = text.replace('\n\n', '\n\n<|endoftext|>')
+    # encode the text
+    tokens = encode(text)
     # let's take the first 32,768 tokens as the validation split (~10%)
     val_tokens = tokens[:32768]
     train_tokens = tokens[32768:]
     # save to file
     val_filename = os.path.join(DATA_CACHE_DIR, "tiny_shakespeare_val.bin")
     train_filename = os.path.join(DATA_CACHE_DIR, "tiny_shakespeare_train.bin")
-    write_datafile(val_filename, val_tokens, model_desc)
-    write_datafile(train_filename, train_tokens, model_desc)
+    write_datafile(val_filename, val_tokens)
+    write_datafile(train_filename, train_tokens)
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Tiny Shakespeare dataset preprocessing")
-    parser.add_argument("-m", "--model_desc", type=str, default="gpt-2", choices=["gpt-2", "llama-3"], help="Model type, gpt-2|llama-3")
-    args = parser.parse_args()
     download()
-    tokenize(args.model_desc)
+    tokenize()

dev/data/tinystories.py

@@ -1,45 +1,38 @@
 """
 Downloads and tokenizes the TinyStories dataset.
 - The download is from HuggingFace datasets.
-- The tokenization is using either GPT-2 or LLaMA 3 tokenizer.
+- The tokenization is GPT-2 tokenizer with tiktoken
 
 The output is written to a newly created tinystories/ folder.
 The script prints:
 
-For GPT-2:
-Number of shards: 50
 Tokenizing val split...
-writing 19,043,638 tokens to tinystories/TinyStories_val.bin
+Saved 19043638 tokens to tinystories/TinyStories_val.bin
 Tokenizing train split...
-writing 925,653,391 tokens to tinystories/TinyStories_train.bin
+Saved 925653391 tokens to tinystories/TinyStories_train.bin
 
-For LLaMA 3:
-Number of shards: 50
-Tokenizing val split...
-writing 18,660,516 tokens to tinystories/TinyStories_val.bin
-Tokenizing train split...
-writing 907,021,844 tokens to tinystories/TinyStories_train.bin
-
-And runs in few minutes two depending on your internet
+And runs in 1-2 minutes two depending on your internet
 connection and computer. The .bin files are raw byte
-streams of uint16 (gpt-2) or uint32 (llama) numbers indicating the token ids.
+streams of int32 numbers indicating the token ids.
 """
 
-import argparse
 import os
 import glob
 import json
 import random
+import requests
+from tqdm import tqdm
 from concurrent.futures import ProcessPoolExecutor, as_completed
-
 import tiktoken
-from transformers import AutoTokenizer
-
+import numpy as np
 from data_common import download_file, write_datafile
 
 # -----------------------------------------------------------------------------
 DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), "tinystories")
 
+enc = tiktoken.get_encoding("gpt2")
+encode = lambda s: enc.encode_ordinary(s)
+
 def download():
     """Downloads the TinyStories dataset to DATA_CACHE_DIR"""
     os.makedirs(DATA_CACHE_DIR, exist_ok=True)
@@ -70,20 +63,10 @@ def download():
     #     data = json.load(f)
     # print(f"Example story:\n{data[0]}")
 
-def process_shard(shard_index, shard_filename, model_desc):
-    if model_desc == "gpt-2":
-        enc = tiktoken.get_encoding("gpt2")
-        encode = lambda s: enc.encode_ordinary(s)
-        eot = enc._special_tokens['<|endoftext|>'] # end of text token
-    elif model_desc == "llama-3":
-        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
-        encode = lambda s: tokenizer.encode(s, add_special_tokens=False, verbose=False, split_special_tokens=True)
-        eot = tokenizer.encode('')[0] # by default the tokenizer adds the EOT token (128000)
-    else:
-        raise ValueError(f"unknown model descriptor {model_desc}")
-
+def process_shard(shard_index, shard_filename):
     with open(shard_filename, "r") as f:
         data = json.load(f)
+    eot = enc._special_tokens['<|endoftext|>'] # end of text token
     rng = random.Random(1337 + shard_index)
     rng.shuffle(data)
     all_tokens = []
@@ -95,7 +78,7 @@ def process_shard(shard_index, shard_filename, model_desc):
         all_tokens.extend(tokens)
     return all_tokens
 
-def tokenize(model_desc):
+def tokenize():
     # shard 0 will be the val split, rest is train
     data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
     shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
@@ -106,17 +89,20 @@ def tokenize(model_desc):
         print(f"Tokenizing {split_name} split...")
         all_tokens = []
         with ProcessPoolExecutor() as executor:
-            futures = [executor.submit(process_shard, shard_index, shard_filename, model_desc)
+            futures = [executor.submit(process_shard, shard_index, shard_filename)
                        for shard_index, shard_filename in enumerate(split_shards)]
             for future in as_completed(futures):
                 all_tokens.extend(future.result())
 
         split_filename = os.path.join(DATA_CACHE_DIR, f"TinyStories_{split_name}.bin")
-        write_datafile(split_filename, all_tokens, model_desc)
+        write_datafile(split_filename, all_tokens)
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Tiny Stories dataset preprocessing")
-    parser.add_argument("-m", "--model_desc", type=str, default="gpt-2", choices=["gpt-2", "llama-3"], help="Model type, gpt-2|llama-3")
-    args = parser.parse_args()
     download()
-    tokenize(args.model_desc)
+    tokenize()
+
+    # Prints:
+    # Tokenizing val split...
+    # Saved 19043638 tokens to data/TinyStories_val.bin
+    # Tokenizing train split...
+    # Saved 925653391 tokens to data/TinyStories_train.bin

dev/download_starter_pack.sh

@@ -4,16 +4,14 @@
 SCRIPT_DIR=$(dirname "$(realpath "$0")")
 
 # Base URL
-BASE_URL="https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/"
+BASE_URL="https://huggingface.co/datasets/chrisdryden/llmcDatasets/resolve/main/"
 
 # Directory paths based on script location
 SAVE_DIR_PARENT="$SCRIPT_DIR/.."
 SAVE_DIR_TINY="$SCRIPT_DIR/data/tinyshakespeare"
-SAVE_DIR_HELLA="$SCRIPT_DIR/data/hellaswag"
 
 # Create the directories if they don't exist
 mkdir -p "$SAVE_DIR_TINY"
-mkdir -p "$SAVE_DIR_HELLA"
 
 # Files to download
 FILES=(
@@ -23,7 +21,6 @@ FILES=(
     "gpt2_tokenizer.bin"
     "tiny_shakespeare_train.bin"
     "tiny_shakespeare_val.bin"
-    "hellaswag_val.bin"
 )
 
 # Function to download files to the appropriate directory
@@ -35,8 +32,6 @@ download_file() {
     # Determine the save directory based on the file name
     if [[ "$FILE_NAME" == tiny_shakespeare* ]]; then
         FILE_PATH="${SAVE_DIR_TINY}/${FILE_NAME}"
-    elif [[ "$FILE_NAME" == hellaswag* ]]; then
-        FILE_PATH="${SAVE_DIR_HELLA}/${FILE_NAME}"
     else
         FILE_PATH="${SAVE_DIR_PARENT}/${FILE_NAME}"
     fi

dev/loss_checker_ci.py

@@ -1,66 +0,0 @@
-# Description: A script to compare numbers in a file with fixed values and check for accuracy within a specified percent difference.
-# Usage: python loss_checker_ci.py -f <file_path> -s <col_start> -e <col_end> -a <percent_accuracy>
-# Example: python dev/loss_checker_ci.py -f train_gpt2cu_fp32_precision.txt -s 20 -e 28 -a 10.0
-import sys
-import argparse
-
-def read_numbers_from_file(file_path, col_start, col_end):
-    try:
-        numbers = []
-        with open(file_path, 'r') as file:
-            lines = file.readlines()
-            start_index = None
-            for i, line in enumerate(lines):
-                if "step    1/10" in line:
-                    start_index = i
-                    break
-
-            if start_index is None:
-                print("Error: Could not find the string 'step    1/10' in the file.")
-                return None
-
-            # Read 10 rows starting from the identified start row
-            for line in lines[start_index:start_index + 10]:
-                # Extracting the specified columns
-                number = float(line[col_start:col_end].strip())
-                numbers.append(number)
-        return numbers
-    except Exception as e:
-        print(f"Error reading the file: {e}")
-        return None
-
-def compare_numbers(read_values, fixed_values, percent_accuracy):
-    for i in range(len(read_values)):
-        read_value = read_values[i]
-        fixed_value = fixed_values[i]
-        percent_difference = ((read_value - fixed_value) / fixed_value) * 100
-        print(f"Fixed Value: {fixed_value}, Read Value: {read_value}, Percent Difference: {percent_difference:.2f}%")
-        if abs(percent_difference) > percent_accuracy:
-            print(f"Error: Percent difference {percent_difference:.2f}% exceeds the allowed accuracy of {percent_accuracy}%")
-            return 1
-    print("Success: All values are within the allowed accuracy.")
-    return 0
-
-def main():
-    parser = argparse.ArgumentParser(description='Compare numbers in a file with fixed values.')
-    parser.add_argument('-f', '--file', required=True, help='Path to the input file')
-    parser.add_argument('-s', '--col_start', type=int, required=True, help='Starting column index (0-based)')
-    parser.add_argument('-e', '--col_end', type=int, required=True, help='Ending column index (0-based)')
-    parser.add_argument('-a', '--percent_accuracy', type=float, required=True, help='Allowed percent accuracy for comparison')
-
-    args = parser.parse_args()
-
-    # Read numbers from file
-    read_values = read_numbers_from_file(args.file, args.col_start, args.col_end)
-    if read_values is None:
-        return 1
-
-    # Use values from test_gpt2.cu for fp32 precision
-    fixed_values = [5.270009,4.060681,3.320085,2.717550,2.181066,1.653923,1.168050,0.736873,0.401021,0.187493];
-
-    # Compare the numbers and check accuracy
-    result = compare_numbers(read_values, fixed_values, args.percent_accuracy)
-    return result
-
-if __name__ == "__main__":
-    sys.exit(main())

dev/test/Makefile

@@ -1,166 +0,0 @@
-CC ?= gcc
-# example: make test_dataloader TEST_CFLAGS=-fsanitize=address -fno-omit-frame-pointer 
-CFLAGS = -Ofast -Wno-unused-result -Wno-ignored-pragmas -Wno-unknown-attributes -g
-CFLAGS += $(TEST_CFLAGS)
-LDFLAGS =
-LDLIBS = -lm
-INCLUDES =
-CFLAGS_COND = -march=native
-
-# Find nvcc
-SHELL_UNAME = $(shell uname)
-REMOVE_FILES = rm -f
-OUTPUT_FILE = -o $@
-CUDA_OUTPUT_FILE = -o $@
-
-# NVCC flags
-# -t=0 is short for --threads, 0 = number of CPUs on the machine
-NVCC_FLAGS = -O3 -t=0 --use_fast_math -std=c++17
-NVCC_LDFLAGS = -lcublas -lcublasLt
-NVCC_INCLUDES =
-NVCC_LDLIBS =
-NVCC_CUDNN =
-# By default we don't build with cudnn because it blows up compile time from a few seconds to ~minute
-USE_CUDNN ?= 0
-
-# We will place .o files in the `build` directory (create it if it doesn't exist)
-BUILD_DIR = build
-$(shell mkdir -p $(BUILD_DIR))
-REMOVE_BUILD_OBJECT_FILES := rm -f $(BUILD_DIR)/*.o
-
-# Function to check if a file exists in the PATH
-define file_exists_in_path
-  $(which $(1) 2>/dev/null)
-endef
-
-ifneq ($(CI),true) # if not in CI, then use the GPU query
-  ifndef GPU_COMPUTE_CAPABILITY # set to defaults if: make GPU_COMPUTE_CAPABILITY=
-    ifneq ($(call file_exists_in_path, __nvcc_device_query),)
-      GPU_COMPUTE_CAPABILITY = $(shell __nvcc_device_query)
-      GPU_COMPUTE_CAPABILITY := $(strip $(GPU_COMPUTE_CAPABILITY))
-    endif
-  endif
-endif
-
-# set to defaults if - make GPU_COMPUTE_CAPABILITY= otherwise use the compute capability detected above
-ifneq ($(GPU_COMPUTE_CAPABILITY),)
-  NVCC_FLAGS += --generate-code arch=compute_$(GPU_COMPUTE_CAPABILITY),code=[compute_$(GPU_COMPUTE_CAPABILITY),sm_$(GPU_COMPUTE_CAPABILITY)]
-endif
-
-# autodect a lot of various supports on current platform
-$(info ---------------------------------------------)
-
-NVCC := $(shell which nvcc 2>/dev/null)
-
-# Check and include cudnn if available
-# You can override the path to cudnn frontend by setting CUDNN_FRONTEND_PATH on the make command line
-# By default, we look for it in HOME/cudnn-frontend/include and ./cudnn-frontend/include
-# Refer to the README for cuDNN install instructions
-ifeq ($(USE_CUDNN), 1)
-  ifeq ($(shell [ -d $(HOME)/cudnn-frontend/include ] && echo "exists"), exists)
-    $(info ✓ cuDNN found, will run with flash-attention)
-    CUDNN_FRONTEND_PATH ?= $(HOME)/cudnn-frontend/include
-  else ifeq ($(shell [ -d cudnn-frontend/include ] && echo "exists"), exists)
-    $(info ✓ cuDNN found, will run with flash-attention)
-    CUDNN_FRONTEND_PATH ?= cudnn-frontend/include
-  else
-    $(error ✗ cuDNN not found. See the README for install instructions and the Makefile for hard-coded paths)
-  endif
-  NVCC_INCLUDES += -I$(CUDNN_FRONTEND_PATH)
-  NVCC_LDFLAGS += -lcudnn
-  NVCC_FLAGS += -DENABLE_CUDNN
-  NVCC_CUDNN = $(BUILD_DIR)/cudnn_att.o
-else
-  $(info → cuDNN is manually disabled by default, run make with `USE_CUDNN=1` to try to enable)
-endif
-
-# Check if OpenMP is available
-# This is done by attempting to compile an empty file with OpenMP flags
-# OpenMP makes the code a lot faster so I advise installing it
-# e.g. on MacOS: brew install libomp
-# e.g. on Ubuntu: sudo apt-get install libomp-dev
-# later, run the program by prepending the number of threads, e.g.: OMP_NUM_THREADS=8 ./gpt2
-# First, check if NO_OMP is set to 1, if not, proceed with the OpenMP checks
-ifeq ($(NO_OMP), 1)
-  $(info OpenMP is manually disabled)
-else
-  ifneq ($(OS), Windows_NT)
-    # Check for OpenMP support in GCC or Clang on Linux
-    ifeq ($(shell echo | $(CC) -fopenmp -x c -E - > /dev/null 2>&1; echo $$?), 0)
-      CFLAGS += -fopenmp -DOMP
-      LDLIBS += -lgomp
-      $(info ✓ OpenMP found)
-    else
-      $(info ✗ OpenMP not found)
-    endif
-  endif
-endif
-
-# Check if OpenMPI and NCCL are available, include them if so, for multi-GPU training
-ifeq ($(NO_MULTI_GPU), 1)
-  $(info → Multi-GPU (OpenMPI + NCCL) is manually disabled)
-else
-  ifeq ($(shell [ -d /usr/lib/x86_64-linux-gnu/openmpi/lib/ ] && [ -d /usr/lib/x86_64-linux-gnu/openmpi/include/ ] && echo "exists"), exists)
-    $(info ✓ OpenMPI found, OK to train with multiple GPUs)
-    NVCC_INCLUDES += -I/usr/lib/x86_64-linux-gnu/openmpi/include
-    NVCC_LDFLAGS += -L/usr/lib/x86_64-linux-gnu/openmpi/lib/
-    NVCC_LDLIBS += -lmpi -lnccl
-    NVCC_FLAGS += -DMULTI_GPU
-  else
-    $(info ✗ OpenMPI is not found, disabling multi-GPU support)
-    $(info ---> On Linux you can try install OpenMPI with `sudo apt install openmpi-bin openmpi-doc libopenmpi-dev`)
-  endif
-endif
-
-# Precision settings, default to bf16 but ability to override
-ifeq ($(MAKECMDGOALS), clean)
-  PRECISION=BF16 
-endif
-
-VALID_PRECISIONS := FP32 FP16 BF16
-ifeq ($(filter $(PRECISION),$(VALID_PRECISIONS)),)
-  $(error Invalid precision $(PRECISION), valid precisions are $(VALID_PRECISIONS))
-endif
-ifeq ($(PRECISION), FP32)
-  PFLAGS = -DENABLE_FP32
-else ifeq ($(PRECISION), FP16)
-  PFLAGS = -DENABLE_FP16
-else
-  PFLAGS = -DENABLE_BF16
-endif
-
-# PHONY means these targets will always be executed
-.PHONY: all clean
-
-# Add targets
-TARGETS = test_dataloader
-
-# Dependency files
-test_dataloader_dependencies = test_dataloader.d
-HEADER_DEPENDENCIES = $(test_dataloader_dependencies)
-
-# Conditional inclusion of CUDA targets
-ifeq ($(NVCC),)
-    $(info ✗ nvcc not found, skipping GPU/CUDA builds)
-else
-    $(info ✓ nvcc found, including GPU/CUDA support)
-    TARGETS += 
-endif
-
-$(info ---------Build Configuration Complete - Build Targets -------------------------)
-
-all: $(TARGETS)
-
-# Generate dependency files
-%.d: %.c
-	$(CC) $(CFLAGS) -MMD -MP -MF $@ -c $<
-
-# Include the dependency files
--include test_dataloader.d
-
-test_dataloader: test_dataloader.c
-	$(CC) $(CFLAGS) $(INCLUDES) $(LDFLAGS) -MMD -MP $^ $(LDLIBS) $(OUTPUT_FILE)
-
-clean:
-	$(REMOVE_FILES) $(TARGETS) *.d *.o
-	$(REMOVE_BUILD_OBJECT_FILES)

dev/test/test_outlier_detector.c

@@ -1,52 +0,0 @@
-/*
-Tests our OutlierDetector
-
-compile and run as (from dev/test directory)
-gcc -O3 -I../../llmc -o test_outlier_detector test_outlier_detector.c -lm && ./test_outlier_detector
-*/
-
-#include <stdlib.h>
-#include "../../llmc/outlier_detector.h"
-
-int main(void) {
-    OutlierDetector detector;
-    init_detector(&detector);
-
-    srand(1337); // init rng
-
-    // generate OUTLIER_DETECTOR_WINDOW_SIZE * 2 random numbers between -1 and 1
-    for (int i = 0; i < OUTLIER_DETECTOR_WINDOW_SIZE * 2; i++) {
-        double val = (double)rand() / RAND_MAX * 2 - 1;  // Random number between -1 and 1
-        double zscore = update_detector(&detector, val);
-
-        printf("Step %d: Value = %.4f, zscore = %.4f\n", i, val, zscore);
-
-        // check that the first OUTLIER_DETECTOR_WINDOW_SIZE values return nan
-        if (i < OUTLIER_DETECTOR_WINDOW_SIZE) {
-            if (!isnan(zscore)) {
-                printf("Error: Expected nan, got %.4f\n", zscore);
-                return EXIT_FAILURE;
-            }
-        } else {
-            // check that the zscore is within reasonable bounds
-            if (zscore < -3.0 || zscore > 3.0) {
-                printf("Error: Z-score %.4f is outside of expected range\n", zscore);
-                return EXIT_FAILURE;
-            }
-        }
-    }
-
-    // simulate an outlier
-    double outlier = 10.0; // <--- loss spike
-    double zscore = update_detector(&detector, outlier);
-    printf("Outlier Step: Value = %.4f, zscore = %.4f\n", outlier, zscore);
-
-    // check that the z-score here is large
-    if (zscore < 5.0) {
-        printf("Error: Z-score %.4f is not large enough for an outlier\n", zscore);
-        return EXIT_FAILURE;
-    }
-
-    printf("OK\n");
-    return EXIT_SUCCESS;
-}

dev/unistd.h

@@ -4,7 +4,6 @@
 
 #define _CRT_SECURE_NO_WARNINGS
 #define _USE_MATH_DEFINES
-#define WIN32_LEAN_AND_MEAN
 
 #include <stdio.h>
 #include <math.h>
@@ -13,8 +12,6 @@
 #include <string.h>
 #include <direct.h> // for _mkdir and _stat
 #include <io.h> // needed for _access below and _findfirst, _findnext, _findclose
-#pragma comment(lib, "Ws2_32.lib")  // Link Ws2_32.lib for socket functions
-#include <winsock2.h>
 
 #define CLOCK_MONOTONIC 0
 static inline int clock_gettime(int ignore_variable, struct timespec* tv)

llmc/adamw.cuh

@@ -61,16 +61,6 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg
                  );
 }
 
-template <typename Tp>
-__global__ void init_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters,
-                                          ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed) {
-    size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-    if (idx >= num_parameters) { return; }
-    params_memory += blockIdx.y * w_stride; // adjust for layer offset
-    master_params_memory += blockIdx.y * s_stride;
-    stochastic_rounding(master_params_memory[idx], &params_memory[idx], seed);
-}
-
 template <typename Tp, typename Tg>
 void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memory, float* m_memory, float* v_memory, size_t num_parameters,
                   ptrdiff_t w_stride, ptrdiff_t g_stride, ptrdiff_t s_stride,  int num_slices, float learning_rate, float beta1, float beta2, int t, float eps, float weight_decay,
@@ -85,14 +75,4 @@ void adamw_update(Tp* params_memory, float* master_params_memory, Tg* grads_memo
                                                          learning_rate, beta1, beta2, beta1_correction, beta2_correction, eps, weight_decay,
                                                          grad_scale, seed);
     cudaCheck(cudaGetLastError());
-}
-
-template <typename Tp>
-void init_from_master(Tp* params_memory, float* master_params_memory, size_t num_parameters,
-                        ptrdiff_t w_stride, ptrdiff_t s_stride, int num_slices, unsigned int seed, cudaStream_t stream) {
-    int block_size = 512; // must match block size of adamw_update so that RNG also matches
-    int num_blocks = CEIL_DIV(num_parameters, block_size);
-    init_from_master_kernel<<<dim3(num_blocks, num_slices), block_size, 0, stream>>>
-                             (params_memory, master_params_memory, num_parameters, w_stride, s_stride, seed);
-    cudaCheck(cudaGetLastError());
-}
+}

llmc/attention.cuh

@@ -199,6 +199,7 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
     // Note: `inp` is not needed for backward pass, so we re-use it as a scratch buffer.
     // Its contents will be overwritten by this function.
     const int block_size = 256;
+    const float alpha = 1.0f, beta = 0.0f;
 
     // inp is (B, T, 3C) QKV
     // preatt, att are (B, NH, T, T)
@@ -214,8 +215,16 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
     int num_blocks = CEIL_DIV(total_threads, block_size);
     permute_kernel<<<num_blocks, block_size, 0, stream>>>(q, k, v, inp, B, T, NH, HS);
 
-    floatX* preatt = inp; // reuse inp as scratch buffer
-    matmul_cublaslt(preatt, k, q, nullptr, T, T, HS, stream, true, false, B * NH, T * HS, T * HS, T * T);
+    floatX* preatt = inp;
+    cublasCheck(cublasSetStream(cublas_handle, stream));
+    cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
+    cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
+                                     CUBLAS_OP_T, CUBLAS_OP_N,
+                                     T, T, HS, &alpha,
+                                     k, CUBLAS_LOWP, HS, T * HS,
+                                     q, CUBLAS_LOWP, HS, T * HS,
+                                     &beta, preatt, CUBLAS_LOWP, T, T * T,
+                                     B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
 
     // multiply all elements of preatt elementwise by scale
     float scale = 1.f / sqrtf(HS);
@@ -225,7 +234,13 @@ void attention_forward(floatX* out, floatX* qkvr, floatX* att,
     // new approach: first cuBLAS another batched matmul
     floatX* vaccum = inp;
     // y = att @ v # (B, nh, T, T) @ (B, nh, T, hs) -> (B, nh, T, hs)
-    matmul_cublaslt(vaccum, v, att, nullptr, HS, T, T, stream, false, false, B * NH, T * HS, T * T, T * HS);
+    cublasCheck(cublasGemmStridedBatchedEx(cublas_handle,
+                                     CUBLAS_OP_N, CUBLAS_OP_N,
+                                     HS, T, T, &alpha,
+                                     v, CUBLAS_LOWP, HS, T * HS,
+                                     att, CUBLAS_LOWP, T, T * T,
+                                     &beta, vaccum, CUBLAS_LOWP, HS, T * HS,
+                                     B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
 
     // now unpermute
     // y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
@@ -243,6 +258,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* datt, floatX* scrat
     NVTX_RANGE_FN();
     const int block_size = 256;
     const int HS = C / NH; // head size
+    const float alpha = 1.0f, beta = 0.0f;
 
     // unpack convenience pointers into q, k, v
     const floatX *q, *k, *v;
@@ -258,17 +274,27 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* datt, floatX* scrat
     int num_blocks = CEIL_DIV(B * T * C, block_size);
     unpermute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(scratch, dout, B, T, NH, HS);
     // backward into datt
-    matmul_cublaslt(datt, v, scratch, nullptr, T, T, HS, stream, true, false, B * NH, T * HS, T * HS, T * T);
+    cublasCheck(cublasSetStream(cublas_handle, stream));
+    cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
+    cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T, T, HS, &alpha,
+                                           v, CUBLAS_LOWP, HS, T * HS, scratch, CUBLAS_LOWP, HS, T * HS, &beta,
+                                           datt, CUBLAS_LOWP, T, T * T, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
     // backward into dv
-    matmul_cublaslt(dv, scratch, att, nullptr, HS, T, T, stream, false, true, B * NH, T * HS, T * T, T * HS);
+    cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha,
+                                           scratch, CUBLAS_LOWP, HS, T * HS, att, CUBLAS_LOWP, T, T * T, &beta,
+                                           dv, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
     const float scale = 1.0f / sqrtf((float)HS);
     // backward into preatt. this is an in-place operation; datt turns into dpreatt here
     softmax_autoregressive_backward_inplace_kernel<<<dim3(T / 4, B * NH), 256>>>(datt, att, B, T, C, scale);
     const floatX* dpreatt = datt;
     // backward into q
-    matmul_cublaslt(dq, k, dpreatt, nullptr, HS, T, T, stream, false, false, B * NH, T * HS, T * T, T * HS);
+    cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS, T, T, &alpha,
+                                           k, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta,
+                                           dq, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
     // backward into k
-    matmul_cublaslt(dk, q, dpreatt, nullptr, HS, T, T, stream, false, true, B * NH, T * HS, T * T, T * HS);
+    cublasCheck(cublasGemmStridedBatchedEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, HS, T, T, &alpha,
+                                           q, CUBLAS_LOWP, HS, T * HS, dpreatt, CUBLAS_LOWP, T, T * T, &beta,
+                                           dk, CUBLAS_LOWP, HS, T * HS, B * NH, cublas_compute, CUBLAS_GEMM_DEFAULT));
     // backward into inp
     num_blocks = CEIL_DIV(B * NH * T * HS, block_size);
     permute_kernel_backward<<<num_blocks, block_size, 0, stream>>>(dinp, dq, dk, dv, B, T, NH, HS);

llmc/cublas_common.h

@@ -14,10 +14,13 @@ cuBLAS related utils
 // cuBLAS Precision settings
 
 #if defined(ENABLE_FP32)
+typedef float floatX;
 #define CUBLAS_LOWP CUDA_R_32F
 #elif defined(ENABLE_FP16)
+typedef half floatX;
 #define CUBLAS_LOWP CUDA_R_16F
 #else // default to bfloat16
+typedef __nv_bfloat16 floatX;
 #define CUBLAS_LOWP CUDA_R_16BF
 #endif
 
@@ -29,6 +32,7 @@ const size_t cublaslt_workspace_size = 32 * 1024 * 1024;
 void* cublaslt_workspace = NULL;
 cublasComputeType_t cublas_compute = CUBLAS_COMPUTE_32F;
 cublasLtHandle_t cublaslt_handle;
+cublasHandle_t cublas_handle;
 
 // ----------------------------------------------------------------------------
 // Error checking

llmc/cuda_common.h

@@ -8,7 +8,6 @@ Common utilities for CUDA code.
 #include <stdio.h>
 #include <math.h>
 #include <string>
-#include <type_traits>      // std::bool_constant
 #include <cuda_runtime.h>
 #include <nvtx3/nvToolsExt.h>
 #include <nvtx3/nvToolsExtCudaRt.h>
@@ -41,21 +40,17 @@ extern cudaDeviceProp deviceProp;
 // convenience macro for calculating grid/block dimensions for kernels
 #define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
 
-// short-cuts for compile-time boolean values that can be used as function arguments
-constexpr std::bool_constant<true> True;
-constexpr std::bool_constant<true> False;
-
 // ----------------------------------------------------------------------------
 // Error checking
 
-// CUDA error checking. Underscore added so this function can be called directly not just via macro
-inline void cudaCheck_(cudaError_t error, const char *file, int line) {
+// CUDA error checking
+inline 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__))
+#define cudaCheck(err) (cudaCheck(err, __FILE__, __LINE__))
 
 // like cudaFree, but checks for errors _and_ resets the pointer.
 template<class T>

llmc/cuda_utils.cuh

@@ -79,36 +79,6 @@ __device__ void store128cg(ElementType* target, Packed128<ElementType> value) {
 typedef Packed128<float> f128;
 typedef Packed128<floatX> x128;
 
-// ----------------------------------------------------------------------------
-// DType support
-
-// enumerator to indentify the datatype of a tensor.
-enum class DType : uint8_t {
-    FP32, FP16, BF16
-};
-
-// Given a datatype enum, returns the underlying number of bytes
-// for a scalar of that type
-size_t sizeof_dtype(DType type) {
-    switch (type) {
-        case DType::FP32:
-            return sizeof(float);
-        case DType::FP16:
-            return sizeof(half);
-        case DType::BF16:
-            return sizeof(nv_bfloat16);
-        default: // handle or get compiler warning
-            fprintf(stderr, "Unknown datatype\n");
-            exit(EXIT_FAILURE);
-    }
-}
-
-DType dtype_of(float* f) { return DType::FP32; }
-DType dtype_of(nv_bfloat16 * f) { return DType::BF16; }
-DType dtype_of(half * f) { return DType::FP16; }
-
-
-
 // ----------------------------------------------------------------------------
 // Copy, cast functions
 
@@ -183,51 +153,6 @@ __device__ inline float blockReduce(float val, bool final_sync=false, float out_
     return block_val;
 }
 
-// Performs a _deterministic_ sum reduction. determinism is achieved by requiring that only
-// a single block be used.
-template<class Float>
-__global__ void global_sum_single_block_kernel(float* result, const Float* values, size_t count) {
-    assert(gridDim.x == 1);     // only a single block!
-    float thread_sum = 0;
-    for(size_t index = threadIdx.x; index < count; index += blockDim.x) {
-        thread_sum += (float)values[index];
-    }
-
-    float reduction = blockReduce<warpReduceSum>(thread_sum, true);
-    if(threadIdx.x == 0) {
-        *result = reduction;
-    }
-}
-
-template<class Float>
-void global_sum_deterministic(float* result, const Float* values, int count, cudaStream_t stream) {
-    global_sum_single_block_kernel<<<1, 1024, 0, stream>>>(result, values, count);
-    cudaCheck(cudaGetLastError());
-}
-
-// ----------------------------------------------------------------------------
-// memory management
-
-// allocate memory, preferrably on the device
-// returns a status code. 0 = OK, 1 = fell back to managed memory
-int cudaMallocConditionallyManaged(void** out, size_t bytes, const char *file, int line) {
-    // try to allocate
-    cudaError_t err = cudaMalloc(out, bytes);
-    if(err == cudaErrorMemoryAllocation) {
-        // if we OOM, fallback to a managed allocation. slower but at least won't crash.
-        cudaGetLastError(); // reset the error before the next API call
-        cudaCheck_(cudaMallocManaged(out, bytes), file, line);
-        cudaCheck_(cudaMemAdvise(*out, bytes, cudaMemAdviseSetPreferredLocation, cudaCpuDeviceId), file, line);
-        return 1;
-    } else {
-        cudaCheck_(err, file, line);
-        return 0;
-    }
-}
-
-#define cudaMallocConditionallyManaged(out, bytes)\
-(cudaMallocConditionallyManaged((void**)out, bytes, __FILE__, __LINE__))
-
 // ----------------------------------------------------------------------------
 // Random Number Generation used in Stochastic Rounding
 
@@ -235,14 +160,14 @@ int cudaMallocConditionallyManaged(void** out, size_t bytes, const char *file, i
 // This gives us a random number from threadIdx/blockIdx + a single seed for the entire GPU
 // todo - possibly overkill and we don't need such high quality random numbers? (tbd)
 // http://eiserloh.net/noise/SquirrelNoise5.hpp
-__device__ __host__ constexpr unsigned int SquirrelNoise5(unsigned int positionX, unsigned int seed)
+__device__ __host__ constexpr unsigned int SquirrelNoise5(int positionX, unsigned int seed)
 {
     constexpr unsigned int SQ5_BIT_NOISE1 = 0xd2a80a3f;	// 11010010101010000000101000111111
     constexpr unsigned int SQ5_BIT_NOISE2 = 0xa884f197;	// 10101000100001001111000110010111
     constexpr unsigned int SQ5_BIT_NOISE3 = 0x6C736F4B; // 01101100011100110110111101001011
     constexpr unsigned int SQ5_BIT_NOISE4 = 0xB79F3ABB;	// 10110111100111110011101010111011
     constexpr unsigned int SQ5_BIT_NOISE5 = 0x1b56c4f5;	// 00011011010101101100010011110101
-    unsigned int mangledBits = positionX;
+    unsigned int mangledBits = (unsigned int) positionX;
     mangledBits *= SQ5_BIT_NOISE1;
     mangledBits += seed;
     mangledBits ^= (mangledBits >> 9);
@@ -258,11 +183,8 @@ __device__ __host__ constexpr unsigned int SquirrelNoise5(unsigned int positionX
 }
 __device__ __host__ constexpr unsigned int Get2dNoiseUint(int indexX, int indexY, unsigned int seed)
 {
-    constexpr unsigned int PRIME_NUMBER = 198491317u; // Large prime number with non-boring bits
-    unsigned int x = static_cast<unsigned int>(indexX);
-    unsigned int y = static_cast<unsigned int>(indexY);
-
-    return SquirrelNoise5(x + (PRIME_NUMBER * y), seed);
+    constexpr int PRIME_NUMBER = 198491317; // Large prime number with non-boring bits
+    return SquirrelNoise5(indexX + (PRIME_NUMBER * indexY), seed);
 }
 
 // stochastic rounding built on top of Squirel Noise above (with seed updated per step via xorshift)

llmc/cudnn_att.cpp

@@ -3,7 +3,6 @@
 // TODO this currently duplicates some of the utilities from the main file
 
 #define NOMINMAX
-#include <unistd.h>
 #include "cudnn_att.h"
 #include <cudnn_frontend.h>
 
@@ -184,9 +183,6 @@ auto lookup_cache_or_build_graph_bwd(int B, int NH, int T, int HS) {
                             .set_uid(Attn_scale_UID)
                             .set_data_type(fe::DataType_t::FLOAT));
     auto sdpa_backward_options = fe::graph::SDPA_backward_attributes().set_name("flash_attention_backward")
-#if CUDNN_FRONTEND_MAJOR_VERSION > 1 || CUDNN_FRONTEND_MINOR_VERSION >= 5
-                            .set_deterministic_algorithm(true) // 1.5+ needs this for determinism
-#endif
                             .set_causal_mask(true)
                             .set_attn_scale(attn_scale);
 

llmc/dataloader.h

@@ -103,8 +103,8 @@ void prepare_intra_shard_indices_(DataLoader *loader) {
         free(loader->intra_shard_indices);
     }
     loader->intra_shard_indices = (int*)mallocCheck(loader->shard_num_samples * sizeof(int));
-    init_identity_permutation(loader->intra_shard_indices, (int) loader->shard_num_samples);
-    random_permutation(loader->intra_shard_indices, (int) loader->shard_num_samples, &loader->shuffle_rng);
+    init_identity_permutation(loader->intra_shard_indices, loader->shard_num_samples);
+    random_permutation(loader->intra_shard_indices, loader->shard_num_samples, &loader->shuffle_rng);
 }
 
 void dataloader_reset(DataLoader *loader) {
@@ -112,10 +112,10 @@ void dataloader_reset(DataLoader *loader) {
     loader->current_sample_idx = 0;
 
     if (loader->should_shuffle) {  // shuffle the shards
-        random_permutation(loader->shard_indices, (int) loader->glob_result.gl_pathc, &loader->shuffle_rng);
+        random_permutation(loader->shard_indices, loader->glob_result.gl_pathc, &loader->shuffle_rng);
     }
 
-    dataloader_load_shard_(loader, (int) loader->current_shard_idx);
+    dataloader_load_shard_(loader, loader->current_shard_idx);
 
     if (loader->should_shuffle) {
         prepare_intra_shard_indices_(loader);
@@ -132,7 +132,7 @@ void dataloader_advance_(DataLoader *loader) {
     // advance the loader by loading the next data shard and resetting the position
     loader->current_shard_idx = (loader->current_shard_idx + 1) % loader->glob_result.gl_pathc;
     loader->current_sample_idx = 0;
-    dataloader_load_shard_(loader, (int) loader->current_shard_idx);
+    dataloader_load_shard_(loader, loader->current_shard_idx);
 
     if (loader->should_shuffle) {
         prepare_intra_shard_indices_(loader);
@@ -172,7 +172,7 @@ void dataloader_init(DataLoader *loader,
         manual_seed(&shuffle_rng, 42 + process_rank);
         loader->shuffle_rng = shuffle_rng;
         loader->shard_indices = (int*)mallocCheck(loader->glob_result.gl_pathc * sizeof(int));
-        init_identity_permutation(loader->shard_indices, (int) loader->glob_result.gl_pathc);
+        init_identity_permutation(loader->shard_indices, loader->glob_result.gl_pathc);
         loader->intra_shard_indices = NULL;  // dynamically allocated allowing different shard sizes
     }
 
@@ -183,7 +183,7 @@ void dataloader_init(DataLoader *loader,
         int64_t shard_ntok = dataloader_load_shard_(loader, shard_index);
         // we need at least one batch/shard, the way things are written right now.
         // can be relaxed a lot later.
-        assert(shard_ntok >= (int64_t) (num_processes * B * T + 1));
+        assert(shard_ntok >= num_processes * B * T + 1);
         ntok_total += shard_ntok;
     }
     // debugging prints
@@ -210,7 +210,7 @@ void dataloader_load_batch(DataLoader* loader) {
     size_t B = loader->B;
     size_t T = loader->T;
     // read B*T+1 uint16_t tokens from the file into buffer
-    fseekCheck(loader->tokens_file, (int) current_offset, SEEK_SET);
+    fseekCheck(loader->tokens_file, current_offset, SEEK_SET);
     freadCheck(loader->buffer, sizeof(uint16_t), B*T+1, loader->tokens_file);
     // decode the buffer into inputs and targets (cast to int)
     for (int i = 0; i < B*T; i++) {
@@ -233,7 +233,7 @@ void dataloader_resume(DataLoader *loader, size_t current_shard_idx, size_t curr
     // used during model resumption (-y 1) flag
     loader->current_shard_idx = current_shard_idx;
     loader->current_sample_idx = current_sample_idx;
-    dataloader_load_shard_(loader, (int) loader->current_shard_idx);
+    dataloader_load_shard_(loader, loader->current_shard_idx);
 }
 
 void dataloader_free(DataLoader *loader) {
@@ -301,13 +301,7 @@ void evalloader_reset(EvalLoader *loader) {
     // then process 0 should start at 0, process 1 at N/4, process 2 at N/2, etc.
     // determine how much work there is for all processes
     int examples_per_process = CEIL_DIV(loader->num_examples, loader->num_processes);
-    int can_fit_examples = (int) (loader->B / ASSUMED_NUM_COMPLETIONS);
-    if (can_fit_examples == 0) {
-        // this could be fixed in the future, but for now keeping it simple and throw error when B too low
-        printf("HellaSwag EvalLoader: batch size %zu is < %d\n", loader->B, ASSUMED_NUM_COMPLETIONS);
-        printf("---> HINT: Disable HellaSwag eval with -h 0, or increase batch size with -b\n");
-        exit(EXIT_FAILURE);
-    }
+    int can_fit_examples = loader->B / ASSUMED_NUM_COMPLETIONS;
     loader->num_batches = CEIL_DIV(examples_per_process, can_fit_examples);
     // determine the start and end example indices for this process
     loader->start_example_index = examples_per_process * loader->process_rank;
@@ -319,7 +313,7 @@ void evalloader_reset(EvalLoader *loader) {
     // now seek through the file to the start of that example
     // utilize <EXAMPLE_BYTES> for efficiency
     int64_t header_bytes = HEADER_SIZE * sizeof(int);
-    fseekCheck(loader->eval_file, (int) header_bytes, SEEK_SET);
+    fseekCheck(loader->eval_file, header_bytes, SEEK_SET);
     for (int i = 0; i < loader->start_example_index; i++) {
         uint16_t example_header[3];
         // read 3 uint16_t values: <START_EXAMPLE>, <EXAMPLE_BYTES>, <EXAMPLE_INDEX>
@@ -331,7 +325,7 @@ void evalloader_reset(EvalLoader *loader) {
         // skip to the next example, keeping in mind that we already read the header
         size_t remaining_bytes = example_header[1] - sizeof(uint16_t) * 3;
         assert(remaining_bytes > 0); // we expect some bytes in the example
-        fseekCheck(loader->eval_file, (int) remaining_bytes, SEEK_CUR);
+        fseekCheck(loader->eval_file, remaining_bytes, SEEK_CUR);
     }
     // now we are at the start of the example we want to start at, pointing at <START_EXAMPLE>
     loader->current_example_index = loader->start_example_index;
@@ -366,7 +360,7 @@ void evalloader_init(EvalLoader *loader,
     assert(longest_example_bytes > 0 && longest_example_bytes < (1+ASSUMED_NUM_COMPLETIONS)*T*2);
 
     // allocate all the space we'll need
-    int can_fit_examples = (int) (B / ASSUMED_NUM_COMPLETIONS);
+    int can_fit_examples = B / ASSUMED_NUM_COMPLETIONS;
     loader->buffer = (uint16_t*)mallocCheck(longest_example_bytes);
     loader->inputs = (int*)calloc(B * T, sizeof(int));
     loader->targets = (int*)calloc(B * T, sizeof(int));
@@ -399,7 +393,7 @@ void evalloader_next_example_(EvalLoader *loader, int example_batch_index) {
     freadCheck(loader->buffer, sizeof(char), example_bytes, loader->eval_file);
     // process the example label
     int label = (int)loader->buffer[0];
-    int can_fit_examples = (int) (loader->B / ASSUMED_NUM_COMPLETIONS);
+    int can_fit_examples = loader->B / ASSUMED_NUM_COMPLETIONS;
     assert(label >= 0 && label < ASSUMED_NUM_COMPLETIONS); // we expect the label to be in [0, 4) for right now
     assert(example_batch_index >= 0 && example_batch_index < can_fit_examples);
     loader->label[example_batch_index] = label; // store for output
@@ -456,7 +450,7 @@ void evalloader_next_batch(EvalLoader *loader) {
     // we have a batch dimension of B, which we want to take full advantage of
     // each example has some number of completions (usually 4)
     // so we want to pack as many examples into rows of B as we can fit
-    int can_fit_examples = (int) (B / ASSUMED_NUM_COMPLETIONS); // how many examples can we fit in the batch?
+    int can_fit_examples = B / ASSUMED_NUM_COMPLETIONS; // how many examples can we fit in the batch?
     for (int i = 0; i < can_fit_examples; i++) {
         if (loader->current_example_index >= loader->end_example_index) {
             break; // this process has exhausted its work, noop from here on
@@ -475,7 +469,7 @@ int evalloader_stat_losses(EvalLoader *loader, float* losses) {
     size_t B = loader->B;
     size_t T = loader->T;
     // iterate the examples in this batch
-    int can_fit_examples = (int) (B / ASSUMED_NUM_COMPLETIONS);
+    int can_fit_examples = B / ASSUMED_NUM_COMPLETIONS;
     for (int i = 0; i < can_fit_examples; i++) {
         float min_loss = 0.0f;
         int min_loss_index = -1;

llmc/fused_classifier.cuh

@@ -65,11 +65,11 @@ __device__ SoftmaxParams prepare_softmax_blockwide3(int64_t idx, const floatX* i
 // will _update_ logits to logit gradients
 // uses template to decide whether to write logits and probs
 // split both loops in "multiple-of-x128-size" and "bounds-checked remainder" parts
-template <bool WriteDLogits = true, bool WriteProbs = false>
+template <bool WriteLogits = true, bool WriteProbs = false>
 __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS)
-    fused_classifier_kernel5(floatX* logits, float* losses, floatX* probs,
+    fused_classifier_kernel5(floatX* logits, floatX* losses, floatX* probs,
                                 const float dloss, const int* targets,
-                                int B, int T, int V, int P, std::bool_constant<WriteDLogits>) {
+                                int B, int T, int V, int P) {
     // note: idx is small enough that it easily fits into 32 bit;
     // by making it a long here, we ensure that any offsets calculated with it (e.g., idx * P)
     // are done is 64 bit
@@ -82,7 +82,7 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS)
     // calculate the probability needed for the loss and update (single-threaded)
     if(threadIdx.x == 0) {
         float prob = expf((float)logits[idx * P + ix] - sp.Offset) * sp.Scale;
-        losses[idx] -= logf(prob);
+        losses[idx] = (floatX)(-logf(prob));
     }
 
     // without this synchronization point we have a race condition:
@@ -106,7 +106,7 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS)
             float indicator = (element == ix) ? 1.0f : 0.0f;
             packed_logits_vec[k] = (floatX)((prob - indicator) * dloss);
         }
-        if (WriteDLogits){
+        if (WriteLogits){
             // reduce cache persistence for the overwritten logits
             // to maximise probability that logits remain in cache between prepare_softmax and here
             store128cs(logits + idx * P + i * x128::size, packed_logits_vec);
@@ -123,7 +123,7 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS)
         float prob = expf((float)logits_vec[i] - sp.Offset) * sp.Scale;
         float indicator = (i == ix) ? 1.0f : 0.0f;
         float dlogit = (prob - indicator) * dloss;
-        if (WriteDLogits){
+        if (WriteLogits){
             __stcs(logits + idx * P + i, (floatX)dlogit);
         }
         if (WriteProbs) {
@@ -136,14 +136,14 @@ __global__ void __launch_bounds__(1024, MAX_1024_THREADS_BLOCKS)
 // kernel launchers
 
 // replaces logits with logit gradients
-template <typename Type, bool WriteDLogits>
-void fused_classifier(Type* logits, float* losses,
+template <typename Type>
+void fused_classifier(Type* logits, Type* losses,
                       const float dloss, const int* targets,
-                      int B, int T, int V, int P, std::bool_constant<WriteDLogits> write_dlogits, cudaStream_t stream) {
+                      int B, int T, int V, int P, cudaStream_t stream) {
     NVTX_RANGE_FN();
     const int block_size = 1024;
     const int N = B * T;
     const int grid_size = N;
-    fused_classifier_kernel5<<<grid_size, block_size, 0, stream>>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P, write_dlogits);
+    fused_classifier_kernel5<<<grid_size, block_size, 0, stream>>>(logits, losses, (floatX*)NULL, dloss, targets, B, T, V, P);
     cudaCheck(cudaGetLastError());
 }

llmc/global_norm.cuh

@@ -87,3 +87,10 @@ void global_norm_squared(float* out, const T* values, size_t count, ptrdiff_t st
     global_norm_squared_kernel<<<dim3(gx, gy), block_size, 0, stream>>>(out, values, count, stride);
     cudaCheck(cudaGetLastError());
 }
+
+void global_norm_squared_aggregate(float* out, int max_num_block_sums, cudaStream_t stream) {
+    assert(max_num_block_sums > 0 && max_num_block_sums < 1024);  // we need to accumulate the block sums in a single block
+    // important to use 1024 here for determinism, otherwise blockreduce might introduce errors
+    global_norm_aggregate_kernel<<<1, 1024, 0, stream>>>(out, max_num_block_sums);
+    cudaCheck(cudaGetLastError());
+}

llmc/layernorm.cuh

@@ -17,7 +17,7 @@ E.g., the layernorms are connected to the residuals so we += in layernorm backwa
 // ----------------------------------------------------------------------------
 // CUDA kernels
 
-__global__ void layernorm_forward_kernel3(floatX* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd,
+__global__ void layernorm_forward_kernel3(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd,
                                     const floatX*  __restrict__ inp, const floatX*  __restrict__ weight,
                                     const floatX* __restrict__ bias, int N, int C) {
     int lane_id = threadIdx.x % WARP_SIZE;
@@ -38,7 +38,7 @@ __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, float* __res
     sum = warpReduceSum(sum);
     float m = sum / C;
     if(lane_id == 0 && mean != nullptr) {
-        __stcs(mean + idx, m);
+        __stcs(mean + idx, (floatX)m);
     }
 
     // rstd
@@ -50,7 +50,7 @@ __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, float* __res
     sum = warpReduceSum(sum);
     float s = rsqrtf(sum / C + 1e-5f);
     if(lane_id == 0 && rstd != nullptr) {
-        __stcs(rstd + idx, s);
+        __stcs(rstd + idx, (floatX)s);
     }
 
     // final normalization and scaling by weight/bias
@@ -64,7 +64,7 @@ __global__ void layernorm_forward_kernel3(floatX* __restrict__ out, float* __res
     }
 }
 
-__global__ void layernorm_forward_kernel6(floatX* __restrict__ out, float* __restrict__ mean, float* __restrict__ rstd,
+__global__ void layernorm_forward_kernel6(floatX* __restrict__ out, floatX* __restrict__ mean, floatX* __restrict__ rstd,
                                     const floatX*  __restrict__ inp, const floatX*  __restrict__ weight,
                                     const floatX* __restrict__ bias, int N, int C) {
     assert(blockDim.x == WARP_SIZE);
@@ -131,15 +131,15 @@ __global__ void layernorm_forward_kernel6(floatX* __restrict__ out, float* __res
     }
     // cache the mean and rstd for the backward pass later
     if(threadIdx.x == 0 && mean != nullptr) {
-        __stcs(mean + idx, m);
+        __stcs(mean + idx, (floatX)m);
     }
     // store the rstd, no need to cache it
     if(threadIdx.x == 0 && rstd != nullptr) {
-        __stcs(rstd + idx, s);
+        __stcs(rstd + idx, (floatX)s);
     }
 }
 
-__global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, float* mean, float* rstd,
+__global__ void fused_residual_forward_kernel5(floatX* residual, floatX* normed, floatX* mean, floatX* rstd,
                                                const floatX* inp1, const floatX* inp2,
                                                const floatX* weight, const floatX* bias,
                                                int N, int C) {
@@ -233,7 +233,7 @@ __global__ void residual_forward_kernel(floatX* out, const floatX* inp1, const f
 __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with only 1024 threads?
     layernorm_backward_kernel10(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch,
                                 const floatX* dout, const floatX* inp, const floatX* weight,
-                                const float* mean, const float* rstd,
+                                const floatX* mean, const floatX* rstd,
                                 int B, int T, int C) {
     int BLOCK_SIZE = blockDim.x;
     int warpsInBlock = BLOCK_SIZE / WARP_SIZE; //number of warps in block
@@ -282,8 +282,8 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with
             }
         }
 
-        const float mean_bt = mean[bt];
-        const float rstd_bt = rstd[bt];
+        const float mean_bt = (float)mean[bt];
+        const float rstd_bt = (float)rstd[bt];
         dnorm_mean = warpReduceSum(dnorm_mean) / C;
         dnorm_norm_mean = warpReduceSum(dnorm_norm_mean) / C * rstd_bt - dnorm_mean * mean_bt * rstd_bt;
 
@@ -430,7 +430,7 @@ __global__ void __launch_bounds__(512, 2) // todo - any warnings on Turing with
 // kernel launchers
 
 // similar to `fused_residual_forward5`
-void layernorm_forward(floatX* out, float* mean, float* rstd,
+void layernorm_forward(floatX* out, floatX* mean, floatX* rstd,
                        floatX* inp, const floatX* weight, const floatX* bias,
                        int B, int T, int C, cudaStream_t stream) {
     NVTX_RANGE_FN();
@@ -444,7 +444,7 @@ void layernorm_forward(floatX* out, float* mean, float* rstd,
     // this may fail, in which case we fall back to the smem free implementation.
     cudaCheck(cudaGetLastError());
     auto status = cudaFuncSetAttribute(layernorm_forward_kernel6, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
-    cudaCheck(cudaGetLastError());
+    cudaGetLastError();
     if (status == cudaSuccess) {
         layernorm_forward_kernel6<<<grid_size, dim3(WARP_SIZE, block_y), smem, stream>>>(out, mean, rstd, inp, weight, bias, N, C);
     } else {
@@ -464,7 +464,7 @@ void residual_forward(floatX* out, const floatX* inp1, const floatX* inp2, int N
     cudaCheck(cudaGetLastError());
 }
 
-void fused_residual_forward5(floatX* residual, floatX* normed, float* mean, float* rstd,
+void fused_residual_forward5(floatX* residual, floatX* normed, floatX* mean, floatX* rstd,
                              const floatX* inp1, const floatX* inp2,
                              const floatX* weight, const floatX* bias,
                              int N, int C, cudaStream_t stream) {
@@ -477,7 +477,7 @@ void fused_residual_forward5(floatX* residual, floatX* normed, float* mean, floa
     // this may fail, in which case we fall back to the smem free implementation.
     cudaCheck(cudaGetLastError());
     auto status = cudaFuncSetAttribute(fused_residual_forward_kernel5, cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
-    cudaCheck(cudaGetLastError());
+    cudaGetLastError();
     if(status == cudaSuccess) {
         fused_residual_forward_kernel5<<<grid_size, dim3(WARP_SIZE, block_y), smem, stream>>>(residual, normed,
                                                                                               mean, rstd, inp1, inp2,
@@ -490,7 +490,7 @@ void fused_residual_forward5(floatX* residual, floatX* normed, float* mean, floa
 }
 
 void layernorm_backward(floatX* dinp, floatX* dweight, floatX* dbias, float* scratch,
-                        const floatX* dout, const floatX* inp, const floatX* weight, const float* mean, const float* rstd,
+                        const floatX* dout, const floatX* inp, const floatX* weight, const floatX* mean, const floatX* rstd,
                         int B, int T, int C, cudaStream_t stream) {
     NVTX_RANGE_FN();
     const int block_size = 512;

llmc/matmul.cuh

@@ -7,8 +7,6 @@ Matrix Multiplication, with help from cuBLASLt
 #include "cuda_common.h"
 #include "cuda_utils.cuh"
 #include "cublas_common.h"
-// GELU can be either fused (cublasLt) or non-fused (gelu.h)
-#include "gelu.cuh"
 
 // ----------------------------------------------------------------------------
 // CUDA kernels
@@ -104,149 +102,83 @@ __global__ void reduce_add_sum_kernel(floatX* dst, const float* src, size_t n, s
 // ----------------------------------------------------------------------------
 // kernel launchers
 
-// Wrapper around cublasLtMatmul that is meant to support everything we need in llm.c
 // https://docs.nvidia.com/cuda/cublas/#cublasltmatmul
-void matmul_cublaslt(floatX* d, const floatX* a, const floatX* b, const floatX* bias,
-                     int m, int n, int k, cudaStream_t stream=0, bool transA=true, bool transB=false,
-                     int batch_count=0, size_t strideA=0, size_t strideB=0, size_t strideOut=0,
-                     bool accumulate=false, floatX* pre_gelu=NULL, bool backward=false)
-{
+void matmul_forward_cublaslt(floatX* out,
+                     floatX* inp, floatX* weight, floatX* bias,
+                     int B, int T, int C, int OC, cudaStream_t stream) {
     NVTX_RANGE_FN();
-    bool has_bias = (bias != NULL);
-    bool has_gelu = (pre_gelu != NULL);
+    int has_bias = (bias != NULL);
 
-    // check alignment (some modes work unaligned but it always best to be aligned for performance)
-    if(((uintptr_t)a % 16) != 0 || ((uintptr_t)b % 16) != 0 || ((uintptr_t)d % 16) != 0 || ((uintptr_t)bias % 16) != 0) {
-        printf("All cuBLASLt pointers must be aligned!\n");
+    // check bias alignment
+    if(((uintptr_t)bias % 16) != 0) {
+        printf("Bias pointer is not aligned (cuBLASLt requirement)!\n");
         exit(EXIT_FAILURE);
     }
 
-    // create the operation descriptor
-    cublasLtMatmulDesc_t operationDesc;
-    cublasCheck(cublasLtMatmulDescCreate(&operationDesc, cublas_compute, CUDA_R_32F));
+    // these need to be in FP16 if and only if alpha/beta are CUBLAS_COMPUTE_16F
+    const float alpha = 1.0f, beta = 0.0f;
 
     int returnedResults = 0;
+    cublasLtMatmulDesc_t operationDesc;
     cublasLtMatmulPreference_t preference;
+    cublasLtMatrixLayout_t weightLayout;
+    cublasLtMatrixLayout_t inputLayout;
+    cublasLtMatrixLayout_t outputLayout;
+    cublasLtMatrixLayout_t biasLayout;
     cublasLtMatmulHeuristicResult_t heuristic;
 
+    // create the operation descriptor
     cublasOperation_t opNoTranspose = CUBLAS_OP_N;
     cublasOperation_t opTranspose = CUBLAS_OP_T;
-    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, (transA)  ? &opTranspose : &opNoTranspose,   sizeof(opTranspose)));
-    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, (transB) ? &opTranspose   : &opNoTranspose, sizeof(opNoTranspose)));
+    cublasLtEpilogue_t epilogueBias = has_bias ? CUBLASLT_EPILOGUE_BIAS : CUBLASLT_EPILOGUE_DEFAULT;
+
+    cublasCheck(cublasLtMatmulDescCreate(&operationDesc, cublas_compute, CUDA_R_32F)); // FP16 if CUBLAS_COMPUTE_16F
+    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &opTranspose, sizeof(opTranspose)));
+    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &opNoTranspose, sizeof(opNoTranspose)));
+    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogueBias, sizeof(epilogueBias)));
+    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));
 
     // define matrix layouts
-    cublasLtMatrixLayout_t ALayout;
-    cublasLtMatrixLayout_t BLayout;
-    cublasLtMatrixLayout_t DLayout;
-    cublasLtMatrixLayout_t CLayout;
-    if (transA) {
-        cublasCheck(cublasLtMatrixLayoutCreate(&ALayout, CUBLAS_LOWP, k, m, k));
-    } else {
-        cublasCheck(cublasLtMatrixLayoutCreate(&ALayout, CUBLAS_LOWP, m, k, m));
-    }
-    if (transB) {
-        cublasCheck(cublasLtMatrixLayoutCreate(&BLayout, CUBLAS_LOWP, n, k, n));
-    } else {
-        cublasCheck(cublasLtMatrixLayoutCreate(&BLayout, CUBLAS_LOWP, k, n, k));
-    }
-    // cuBLASLt requires C in FP8 mode to be BF16 or FP32... (sigh)
-    cublasCheck(cublasLtMatrixLayoutCreate(&CLayout, (sizeof(floatX) == 1) ? CUDA_R_16BF : CUBLAS_LOWP, m, n, m));
-    cublasCheck(cublasLtMatrixLayoutCreate(&DLayout, CUBLAS_LOWP, m, n, m));
-
-    // Strided Batched GEMM (used for non-flash attention, equivalent to cublasGemmStridedBatchedEx)
-    if (batch_count) {
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(ALayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(BLayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(CLayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(DLayout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count, sizeof(batch_count)));
-
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(ALayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideA, sizeof(strideA)));
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(BLayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideB, sizeof(strideB)));
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(CLayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideOut, sizeof(strideOut)));
-        cublasCheck(cublasLtMatrixLayoutSetAttribute(DLayout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &strideOut, sizeof(strideOut)));
-    }
+    cublasCheck(cublasLtMatrixLayoutCreate(&weightLayout, CUBLAS_LOWP, C, OC, C));
+    cublasCheck(cublasLtMatrixLayoutCreate(&inputLayout, CUBLAS_LOWP, C, B*T, C));
+    cublasCheck(cublasLtMatrixLayoutCreate(&outputLayout, CUBLAS_LOWP, OC, B*T, OC));
+    cublasCheck(cublasLtMatrixLayoutCreate(&biasLayout, CUBLAS_LOWP, OC, 1, OC));
 
     // create a preference handle with specified max workspace
     cublasCheck(cublasLtMatmulPreferenceCreate(&preference));
-    cublasCheck(cublasLtMatmulPreferenceSetAttribute(preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
-                                                     &cublaslt_workspace_size, sizeof(cublaslt_workspace_size)));
+    cublasCheck(cublasLtMatmulPreferenceSetAttribute(preference,
+        CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &cublaslt_workspace_size, sizeof(cublaslt_workspace_size)));
 
-    // setup epilogue and associated pointers for bias & gelu
-    cublasLtEpilogue_t epilogue;
-    if (has_gelu) {
-        int64_t gelu_ld = m; // todo - is this affected by anything else?
-        cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &gelu_ld, sizeof(gelu_ld)));
-        cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER, &pre_gelu, sizeof(pre_gelu)));
-        if (backward) {
-            assert(!has_bias); // we shouldn't have any backward matmuls that use both GELU and bias
-            epilogue = CUBLASLT_EPILOGUE_DGELU;
-        } else {
-            epilogue = has_bias ? CUBLASLT_EPILOGUE_GELU_AUX_BIAS : CUBLASLT_EPILOGUE_GELU_AUX;
-        }
-    } else if(has_bias){
-        epilogue = backward ? CUBLASLT_EPILOGUE_BGRADB : CUBLASLT_EPILOGUE_BIAS;
-    } else {
-        epilogue = CUBLASLT_EPILOGUE_DEFAULT;
-    }
-    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue)));
-
-    if (has_bias) {
-        // cuBLASLt requires bias in FP8 mode to be BF16... (sigh)
-        cublasDataType_t bias_data_type = (sizeof(floatX) == 1) ? CUDA_R_16BF : CUBLAS_LOWP; // force BF16 bias for FP8 mode
-        cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE, &bias_data_type, sizeof(bias_data_type)));
-        cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));
-    }
-
-    // set scale type to FP32 (needs to be FP16 if and only if using CUBLAS_COMPUTE_16F, so it's FP32 even for FP8!)
-    cublasDataType_t scale_type = CUDA_R_32F;
-    cublasCheck(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)));
-
-    // find a suitable algorithm (cached internally so shouldn't take much CPU time in practice)
-    cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc, ALayout, BLayout, CLayout, DLayout,
-                                   preference, 1, &heuristic, &returnedResults);
+    // find a suitable algorithm
+    cublasCheck(cublasLtMatmulAlgoGetHeuristic(cublaslt_handle, operationDesc,
+        weightLayout, inputLayout, outputLayout, outputLayout,
+        preference, 1, &heuristic, &returnedResults));
     if (returnedResults == 0) {
-        printf("No cuBLASLt algorithm: m: %d, n: %d, k: %d, bias: %d\n", n, m, k, has_bias);
+        printf("No cuBLASLt algorithm: B: %d, T: %d, C: %d, OC: %d, bias: %d\n", B, T, C, OC, has_bias);
         exit(EXIT_FAILURE);
     }
 
-    // set whether to accumulate (i.e. D += C) or not - note this isn't considered in algorithm selection (?!)
-    const float alpha = 1.0f, beta = accumulate ? 1.0f : 0.0f;
-
     // call the matmul
     cublasCheck(cublasLtMatmul(cublaslt_handle, operationDesc,
-                               &alpha, a, ALayout, b, BLayout, &beta, d, CLayout, d, DLayout,
-                               &heuristic.algo, cublaslt_workspace, cublaslt_workspace_size, stream));
+        &alpha, weight, weightLayout, inp, inputLayout, &beta,
+        out, outputLayout, out, outputLayout, &heuristic.algo,
+        cublaslt_workspace, cublaslt_workspace_size, stream));
 
     // cleanups
     cublasCheck(cublasLtMatmulPreferenceDestroy(preference));
     cublasCheck(cublasLtMatmulDescDestroy(operationDesc));
-    cublasCheck(cublasLtMatrixLayoutDestroy(ALayout));
-    cublasCheck(cublasLtMatrixLayoutDestroy(BLayout));
-    cublasCheck(cublasLtMatrixLayoutDestroy(CLayout));
-    cublasCheck(cublasLtMatrixLayoutDestroy(DLayout));
-    cudaCheck(cudaGetLastError());
-}
-
-// small wrapper around matmul_cublaslt for the forward pass (keeping historical order of arguments)
-void matmul_forward_cublaslt(floatX* out,
-                     floatX* inp, floatX* weight, floatX* bias,
-                     int B, int T, int C, int OC, cudaStream_t stream,
-                     floatX* pre_gelu=NULL, int gelu_fusion=1) {
-    // By default only fuse GELU for H100+ as cuBLAS seems to be inefficient for fused GELU on Ada/Ampere (?)
-    if (gelu_fusion < 1 && pre_gelu) {
-        matmul_cublaslt(pre_gelu, weight, inp, bias, OC, B*T, C, stream, true, false, 0, 0, 0, 0, false, NULL, false);
-        gelu_forward(out, pre_gelu, B*T*OC, stream);
-    } else {
-        matmul_cublaslt(out, weight, inp, bias, OC, B*T, C, stream, true, false, 0, 0, 0, 0, false, pre_gelu, false);
-    }
+    cublasCheck(cublasLtMatrixLayoutDestroy(weightLayout));
+    cublasCheck(cublasLtMatrixLayoutDestroy(inputLayout));
+    cublasCheck(cublasLtMatrixLayoutDestroy(outputLayout));
+    cublasCheck(cublasLtMatrixLayoutDestroy(biasLayout));
 }
 
 void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
                      floatX* dout, floatX* inp, floatX* weight,
                      float* dbias_buffer,
-                     int B, int T, int C, int OC, cudaStream_t stream,
-                     floatX* pre_gelu=NULL, int gelu_fusion=1) {
+                     int B, int T, int C, int OC, cudaStream_t stream) {
     NVTX_RANGE_FN();
+    float one = 1.0f, zero = 0.0f;
 
     // backward to bias, if given, does a +=
     if (dbias != NULL) {
@@ -263,28 +195,26 @@ void matmul_backward(floatX* dinp, floatX* dweight, floatX* dbias,
         // If we have enough OC that we don't need cross-block reductions, we can skip the bias_buffer accumulation
         // and write results directly to the output.
         if(grid_size_y == 1) {
-            matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim, 0, stream>>>(dbias, dout, B, T, OC, False);
+            matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim, 0, stream>>>(dbias, dout, B, T, OC, std::bool_constant<false>{});
             cudaCheck(cudaGetLastError());
         } else {
             // kernel 9 overwrites temp buffer, so no need to memset
-            matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim, 0, stream>>>(dbias_buffer, dout, B, T, OC, True);
+            matmul_backward_bias_kernel9<<<dim3(grid_size_x, grid_size_y), block_dim, 0, stream>>>(dbias_buffer, dout, B, T, OC, std::bool_constant<true>{});
             cudaCheck(cudaGetLastError());
             reduce_add_sum_kernel<<<CEIL_DIV(OC, 256 * f128::size), 256, 0, stream>>>(dbias, dbias_buffer, OC, grid_size_y);
             cudaCheck(cudaGetLastError());
         }
-        dbias = NULL; // prevent dbias calculation from also being fused in matmul_cublaslt below (if we enabled fusion)
     }
 
     // backward to input, uses = in the backward pass (set the gradient)
-    matmul_cublaslt(dinp, weight, dout, NULL, C, B*T, OC, stream, false, false, 0, 0, 0, 0, false,
-                    gelu_fusion >= 2 ? pre_gelu : NULL, true);
-
-    // backward GELU (if it wasn't fused into the matmul above)
-    if (gelu_fusion < 2 && pre_gelu) {
-        gelu_backward_inplace(dinp, pre_gelu, B*T*C, stream);
-    }
-
+    cublasCheck(cublasSetStream(cublas_handle, stream));
+    cublasCheck(cublasSetWorkspace(cublas_handle, cublaslt_workspace, cublaslt_workspace_size));
+    cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, C, B*T, OC, &one,
+                             weight, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &zero,
+                             dinp, CUBLAS_LOWP, C, cublas_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
     // backward to weight, uses += in the backward pass (accumulate the gradient) by setting alpha=one
-    matmul_cublaslt(dweight, inp, dout, NULL /*dbias*/, C, OC, B*T, stream, false, true, 0, 0, 0, 0,
-                    true /* accumulate */, NULL, true);
+    cublasCheck(cublasGemmEx(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, C, OC, B*T, &one,
+                             inp, CUBLAS_LOWP, C, dout, CUBLAS_LOWP, OC, &one,
+                             dweight, CUBLAS_LOWP, C, cublas_compute, CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+    cudaCheck(cudaGetLastError());
 }

llmc/mfu.h

@@ -4,29 +4,12 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
-#if __has_include(<nvml.h>)
-#define USE_NVML 1
-#include <nvml.h>
-#else
-#define USE_NVML 0
-#endif
 
 // tied to enum PrecisionMode, in a future refactor make them the same
 #define MFUH_PRECISION_FP32 0
 #define MFUH_PRECISION_FP16 1
 #define MFUH_PRECISION_BF16 2
 
-#if USE_NVML
-inline void nvml_check(nvmlReturn_t status, const char *file, int line) {
-    if (status != NVML_SUCCESS) {
-        printf("[NVML ERROR] at file %s:%d:\n%s\n", file, line, nvmlErrorString(status));
-        exit(EXIT_FAILURE);
-    }
-};
-#define nvmlCheck(err) (nvml_check(err, __FILE__, __LINE__))
-#endif
-
-
 typedef struct {
     float TF_32;       // tensor-core performance 32 bit
     float BF_16_32;    // bf16 with 32 bit accumulate
@@ -151,94 +134,4 @@ float get_flops_promised(const char* device, int precision_mode) {
     return -1.0f; // ¯\_(ツ)_/¯
 }
 
-struct GPUUtilInfo {
-    unsigned int clock;
-    unsigned int max_clock;
-    unsigned int power;
-    unsigned int power_limit;
-    unsigned int fan;
-    unsigned int temperature;
-    unsigned int temp_slowdown;
-
-    float gpu_utilization;
-    float mem_utilization;
-    const char* throttle_reason;
-};
-
-// lazily initialize nvml and generate a handle to the GPU
-#if USE_NVML
-nvmlDevice_t nvml_get_device() {
-    static bool needs_init = true;
-    static nvmlDevice_t device;
-    if(needs_init) {
-        needs_init = false;
-        nvmlCheck(nvmlInit());
-        nvmlCheck(nvmlDeviceGetHandleByIndex_v2(0, &device));
-    }
-    return device;
-}
-
-// convert throttle reason bitfield into a text reason.
-// this is a lossy conversion; we just want to give some idea of what is happening
-const char* get_throttle_reason(unsigned long long bits) {
-    if(bits & (nvmlClocksThrottleReasonSwPowerCap | nvmlClocksThrottleReasonHwPowerBrakeSlowdown)) {
-        return "power cap";
-    } else if (bits & (nvmlClocksThrottleReasonSwThermalSlowdown | nvmlClocksThrottleReasonHwThermalSlowdown)) {
-        return "thermal cap";
-    } else if (bits & (nvmlClocksThrottleReasonAll)) {
-        return "other cap";
-    } else {
-        return "no cap";
-    }
-}
-
-// gather data for a GPUUtilInfo object
-GPUUtilInfo get_gpu_utilization_info() {
-    GPUUtilInfo info;
-    nvmlDevice_t device = nvml_get_device();
-    // query different infos directly
-    nvmlCheck(nvmlDeviceGetClockInfo(device, NVML_CLOCK_SM, &info.clock));
-    nvmlCheck(nvmlDeviceGetMaxClockInfo(device, NVML_CLOCK_SM, &info.max_clock));
-    nvmlCheck(nvmlDeviceGetPowerManagementLimit(device, &info.power_limit));
-    nvmlCheck(nvmlDeviceGetPowerUsage(device, &info.power));
-    nvmlCheck(nvmlDeviceGetTemperature(device, NVML_TEMPERATURE_GPU, &info.temperature));
-    nvmlCheck(nvmlDeviceGetTemperatureThreshold(device, NVML_TEMPERATURE_THRESHOLD_SLOWDOWN, &info.temp_slowdown));
-    unsigned long long throttle;
-    nvmlCheck(nvmlDeviceGetCurrentClocksThrottleReasons(device, &throttle));
-    info.throttle_reason = get_throttle_reason(throttle);
-    nvmlCheck(nvmlDeviceGetFanSpeed(device, &info.fan));
-
-    // for "utilization", we look at recorded samples. In principle, we could query the driver for how many samples
-    // to request, but then we'd need to dynamically allocate sufficient space. Let's just hard-code a limit of 128,
-    // and have no memory management required
-    constexpr const int BUFFER_LIMIT = 128;
-    nvmlSample_t buffer[BUFFER_LIMIT];
-    nvmlValueType_t v_type;
-    unsigned int sample_count = BUFFER_LIMIT;
-    nvmlCheck(nvmlDeviceGetSamples(device, NVML_GPU_UTILIZATION_SAMPLES, 0, &v_type, &sample_count, buffer));
-    float gpu_utilization = 0.f;
-    for(unsigned i = 0; i < sample_count; ++i) {
-        gpu_utilization += (float)buffer[i].sampleValue.uiVal;
-    }
-    gpu_utilization /= (float)sample_count;
-
-    // sample count may have been modified by the query above; reset back to buffer size
-    sample_count = BUFFER_LIMIT;
-    nvmlCheck(nvmlDeviceGetSamples(device, NVML_MEMORY_UTILIZATION_SAMPLES, 0, &v_type, &sample_count, buffer));
-    float mem_utilization = 0.f;
-    for(unsigned i = 0; i < sample_count; ++i) {
-        mem_utilization += (float)buffer[i].sampleValue.uiVal;
-    }
-    mem_utilization /= (float)sample_count;
-
-    info.gpu_utilization = gpu_utilization;
-    info.mem_utilization = mem_utilization;
-    return info;
-}
-#else
-GPUUtilInfo get_gpu_utilization_info() {
-    fprintf(stderr, "Error: Compiled without nvml support. Cannot perform additional GPU state tracking.");
-    exit(EXIT_FAILURE);
-}
-#endif
 #endif // MFU_H

llmc/outlier_detector.h

@@ -1,70 +0,0 @@
-/*
-Simple OutlierDetector that we can use to monitor the loss and grad norm
-Internally, it keeps track of a window of measurements and each time we
-add a measurement, it returns the z-score of the new value with respect to
-the window of measurements. This can be used to detect outliers in the data.
-
-We use double so that the detector doesn't drift too much, because we
-update the mean and variance with += on each step for efficiency. We could
-reconsider this choice in the future, as the compute cost here is minimal.
-*/
-
-#include <stdio.h>
-#include <math.h>
-
-// use compile-time constant for window size to avoid dynamic memory allocations
-#define OUTLIER_DETECTOR_WINDOW_SIZE 128
-
-typedef struct {
-    double buffer[OUTLIER_DETECTOR_WINDOW_SIZE];
-    int count;
-    int index;
-    double sum;
-    double sum_sq;
-} OutlierDetector;
-
-void init_detector(OutlierDetector *detector) {
-    for (int i = 0; i < OUTLIER_DETECTOR_WINDOW_SIZE; i++) {
-        detector->buffer[i] = 0.0;
-    }
-    detector->count = 0;
-    detector->index = 0;
-    detector->sum = 0.0;
-    detector->sum_sq = 0.0;
-}
-
-double update_detector(OutlierDetector *detector, double new_value) {
-
-    if (detector->count < OUTLIER_DETECTOR_WINDOW_SIZE) {
-        // here we are still building up a window of observations
-        detector->buffer[detector->count] = new_value;
-        detector->sum += new_value;
-        detector->sum_sq += new_value * new_value;
-        detector->count++;
-        return nan(""); // not enough data yet
-
-    } else {
-        // we've filled the window, so now we can start detecting outliers
-
-        // pop the oldest value from the window
-        double old_value = detector->buffer[detector->index];
-        detector->sum -= old_value;
-        detector->sum_sq -= old_value * old_value;
-        // push the new value into the window
-        detector->buffer[detector->index] = new_value;
-        detector->sum += new_value;
-        detector->sum_sq += new_value * new_value;
-        // move the index to the next position
-        detector->index = (detector->index + 1) % OUTLIER_DETECTOR_WINDOW_SIZE;
-        // calculate the z-score of the new value
-        double mean = detector->sum / OUTLIER_DETECTOR_WINDOW_SIZE;
-        double variance = (detector->sum_sq / OUTLIER_DETECTOR_WINDOW_SIZE) - (mean * mean);
-        double std_dev = sqrt(variance);
-        if (std_dev == 0.0) {
-            return 0.0;
-        }
-        double z = (new_value - mean) / std_dev;
-
-        return z;
-    }
-}

llmc/rand.h

@@ -165,13 +165,13 @@ void uniform_(float* data, unsigned int numel, float from, float to, mt19937_sta
 
 // Box-Muller transform: maps uniform random numbers to Gaussian distributed numbers
 // https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
-void normal_fill_16(float* data, float mean, float std) {
-    #define EPSILONE 1e-12f
+void normal_fill_16(float* data, float mean, float std, mt19937_state* state) {
+    #define EPSILONE 1e-12
     for (unsigned int t = 0; t < 8; t++) {
         float u1 = 1 - data[t];
         float u2 = data[t + 8];
         float radius = sqrtf(-2 * logf(u1 + EPSILONE));
-        float theta = (float) (2.0 * M_PI * u2);
+        float theta = 2.0 * M_PI * u2;
         data[t] = (radius * cosf(theta) * std + mean);
         data[t + 8] = (radius * sinf(theta) * std + mean);
     }
@@ -182,7 +182,7 @@ void normal_fill(float* data, unsigned int numel, float mean, float std, mt19937
         data[t] = randfloat32(state);
     }
     for (unsigned int i = 0; i < numel - 15; i += 16) {
-        normal_fill_16(data + i, mean, std);
+        normal_fill_16(data + i, mean, std, state);
     }
     if (numel % 16 != 0) {
         // recompute the last 16 values
@@ -190,12 +190,12 @@ void normal_fill(float* data, unsigned int numel, float mean, float std, mt19937
         for (unsigned int i = 0; i < 16; i++) {
             data[i] = randfloat32(state);
         }
-        normal_fill_16(data, mean, std);
+        normal_fill_16(data, mean, std, state);
     }
 }
 
 void normal_(float* data, unsigned int numel, float mean, float std, mt19937_state* state) {
-    #define EPSILONE 1e-12f
+    #define EPSILONE 1e-12
     if (numel >= 16) {
         normal_fill(data, numel, mean, std, state);
     }
@@ -209,10 +209,10 @@ void normal_(float* data, unsigned int numel, float mean, float std, mt19937_sta
                 continue;
             }
             // for numel < 16 we draw a double (float64)
-            float u1 = (float) randfloat64(state);
-            float u2 = (float) randfloat64(state);
+            float u1 = randfloat64(state);
+            float u2 = randfloat64(state);
             float radius = sqrtf(-2 * logf(1 - u2 + EPSILONE));
-            float theta = (float) (2.0 * M_PI * u1);
+            float theta = 2.0 * M_PI * u1;
             next_double_normal_sample = radius * sinf(theta);
             has_next_double_normal_sample = 1;
             data[t] = (radius * cosf(theta) * std + mean);

llmc/utils.h

@@ -7,7 +7,6 @@
 #ifndef UTILS_H
 #define UTILS_H
 
-#include <unistd.h>
 #include <string.h>
 #include <stdio.h>
 #include <stdlib.h>
@@ -15,7 +14,6 @@
 // implementation of dirent for Windows is in dev/unistd.h
 #ifndef _WIN32
 #include <dirent.h>
-#include <arpa/inet.h>
 #endif
 
 // ----------------------------------------------------------------------------
@@ -75,32 +73,6 @@ extern inline void fclose_check(FILE *fp, const char *file, int line) {
 
 #define fcloseCheck(fp) fclose_check(fp, __FILE__, __LINE__)
 
-extern inline void sclose_check(int sockfd, const char *file, int line) {
-    if (close(sockfd) != 0) {
-        fprintf(stderr, "Error: Failed to close socket at %s:%d\n", file, line);
-        fprintf(stderr, "Error details:\n");
-        fprintf(stderr, "  File: %s\n", file);
-        fprintf(stderr, "  Line: %d\n", line);
-        exit(EXIT_FAILURE);
-    }
-}
-
-#define scloseCheck(sockfd) sclose_check(sockfd, __FILE__, __LINE__)
-
-#ifdef _WIN32
-extern inline void closesocket_check(int sockfd, const char *file, int line) {
-    if (closesocket(sockfd) != 0) {
-        fprintf(stderr, "Error: Failed to close socket at %s:%d\n", file, line);
-        fprintf(stderr, "Error details:\n");
-        fprintf(stderr, "  File: %s\n", file);
-        fprintf(stderr, "  Line: %d\n", line);
-        exit(EXIT_FAILURE);
-    }
-}
-
-#define closesocketCheck(sockfd) closesocket_check(sockfd, __FILE__, __LINE__)
-#endif
-
 extern inline void fseek_check(FILE *fp, long off, int whence, const char *file, int line) {
     if (fseek(fp, off, whence) != 0) {
         fprintf(stderr, "Error: Failed to seek in file at %s:%d\n", file, line);
@@ -155,25 +127,6 @@ extern inline void *malloc_check(size_t size, const char *file, int line) {
 
 #define mallocCheck(size) malloc_check(size, __FILE__, __LINE__)
 
-
-// ----------------------------------------------------------------------------
-// check that all tokens are within range
-extern inline void token_check(const int* tokens, int token_count, int vocab_size, const char *file, int line) {
-    for(int i = 0; i < token_count; i++) {
-        if(!(0 <= tokens[i] && tokens[i] < vocab_size)) {
-            fprintf(stderr, "Error: Token out of vocabulary at %s:%d\n", file, line);
-            fprintf(stderr, "Error details:\n");
-            fprintf(stderr, "  File: %s\n", file);
-            fprintf(stderr, "  Line: %d\n", line);
-            fprintf(stderr, "  Token: %d\n", tokens[i]);
-            fprintf(stderr, "  Position: %d\n", i);
-            fprintf(stderr, "  Vocab: %d\n", vocab_size);
-            exit(EXIT_FAILURE);
-        }
-    }
-}
-#define tokenCheck(tokens, count, vocab) token_check(tokens, count, vocab, __FILE__, __LINE__)
-
 // ----------------------------------------------------------------------------
 // I/O ops
 
@@ -209,15 +162,4 @@ extern inline int find_max_step(const char* output_log_dir) {
     return max_step;
 }
 
-extern inline int ends_with_bin(const char* str) {
-    // checks if str ends with ".bin". could be generalized in the future.
-    if (str == NULL) { return 0; }
-    size_t len = strlen(str);
-    const char* suffix = ".bin";
-    size_t suffix_len = strlen(suffix);
-    if (len < suffix_len) { return 0; }
-    int suffix_matches = strncmp(str + len - suffix_len, suffix, suffix_len) == 0;
-    return suffix_matches;
-}
-
 #endif

llmc/zero.cuh

@@ -5,6 +5,12 @@ Utilities for ZeRO sharding
 #ifndef LLMC_ZERO_CUH
 #define LLMC_ZERO_CUH
 
+#ifdef _WIN32
+#include <winsock2.h>
+#else
+#include <arpa/inet.h>
+#endif
+
 #include <cuda_runtime_api.h>
 #include <stdint.h>
 #include <stdlib.h>
@@ -18,9 +24,6 @@ Utilities for ZeRO sharding
 #endif
 #endif
 
-// defines: fcloseCheck, fwriteCheck, scloseCheck, sclosesocketCheck
-#include "utils.h"
-
 // ----------------------------------------------------------------------------
 // Multi-GPU related
 #ifdef MULTI_GPU
@@ -78,10 +81,6 @@ typedef struct {
 #endif
 } MultiGpuConfig;
 
-// one global variable to hold the multi-GPU configuration for this process
-// inline, so we can include this header multiple times without getting multiple definitions
-inline MultiGpuConfig multi_gpu_config;
-
 #ifdef MULTI_GPU
 
 #ifdef _WIN32
@@ -92,7 +91,7 @@ void send_nccl_id_to_clients_windows(ncclUniqueId *nccl_id, SOCKET client_socket
             WSACleanup();
             exit(EXIT_FAILURE);
         }
-        closesocketCheck(client_sockets[i]);
+        closesocket(client_sockets[i]);
     }
 }
 #else
@@ -102,7 +101,7 @@ void send_nccl_id_to_clients(ncclUniqueId *nccl_id, int client_sockets[], int nu
             printf("Failed to send nccl_id");
             exit(EXIT_FAILURE);
         }
-        scloseCheck(client_sockets[i]);
+        close(client_sockets[i]);
     }
 }
 #endif
@@ -144,7 +143,7 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser
         // Step 3) bind the socket to the address and port
         if (bind(server_socket, (struct sockaddr *)&address, sizeof(address)) == SOCKET_ERROR) {
             printf("Bind failed");
-            closesocketCheck(server_socket);
+            closesocket(server_socket);
             WSACleanup();
             exit(EXIT_FAILURE);
         }
@@ -152,7 +151,7 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser
         // Step 4) MAX_CLIENTS specifies the maximum number of clients that can be queued for this server
         if (listen(server_socket, MAX_CLIENTS) == SOCKET_ERROR) {
             printf("Listen failed");
-            closesocketCheck(server_socket);
+            closesocket(server_socket);
             WSACleanup();
             exit(EXIT_FAILURE);
         }
@@ -162,7 +161,7 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser
         while (num_clients < MAX_CLIENTS) {
             if ((new_socket = accept(server_socket, (struct sockaddr *)&address, &addrlen)) == INVALID_SOCKET) {
                 printf("Accept failed");
-                closesocketCheck(server_socket);
+                closesocket(server_socket);
                 WSACleanup();
                 exit(EXIT_FAILURE);
             }
@@ -174,7 +173,7 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser
         send_nccl_id_to_clients_windows(&nccl_id, client_sockets, num_clients);
         printf("NCCL ID sent to all clients\n");
 
-        closesocketCheck(server_socket);
+        closesocket(server_socket);
     } else {
         int num_connection_attempts = 5;
         int time_to_sleep = 2;
@@ -193,7 +192,7 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser
         serv_addr.sin_port = htons(SERVER_PORT);
         if (inet_pton(AF_INET, server_ip, &serv_addr.sin_addr) <= 0) {
             printf("Invalid address or address not supported");
-            closesocketCheck(client_socket);
+            closesocket(client_socket);
             WSACleanup();
             exit(EXIT_FAILURE);
         }
@@ -203,7 +202,7 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser
             printf("%d Connection failed, retrying in %d seconds\n", result->process_rank, time_to_sleep);
             if (--num_connection_attempts == 0) {
                 printf("Failed to connect to the server\n");
-                closesocketCheck(client_socket);
+                closesocket(client_socket);
                 WSACleanup();
                 exit(EXIT_FAILURE);
             }
@@ -213,13 +212,13 @@ ncclUniqueId get_nccl_id_via_tcp_windows(MultiGpuConfig* result, const char* ser
         // Step 4) receive the NCCL ID from the server
         if (recv(client_socket, (char *)&nccl_id, sizeof(nccl_id), 0) <= 0) {
             printf("Failed to receive nccl_id");
-            closesocketCheck(client_socket);
+            closesocket(client_socket);
             WSACleanup();
             exit(EXIT_FAILURE);
         }
 
         printf("Received NCCL ID\n");
-        closesocketCheck(client_socket);
+        closesocket(client_socket);
     }
 
     WSACleanup();
@@ -288,7 +287,7 @@ ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip)
         send_nccl_id_to_clients(&nccl_id, client_sockets, num_clients);
         printf("NCCL ID sent to all clients\n");
 
-        scloseCheck(server_socket);
+        close(server_socket);
     } else {
         int num_connection_attempts = 5;
         int time_to_sleep = 2;
@@ -326,7 +325,7 @@ ncclUniqueId get_nccl_id_via_tcp(MultiGpuConfig* result, const char* server_ip)
         }
 
         printf("Received NCCL ID\n");
-        scloseCheck(client_socket);
+        close(client_socket);
     }
 
     return nccl_id;
@@ -349,8 +348,8 @@ ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) {
         ncclCheck(ncclGetUniqueId(&nccl_id));
         idFile = fopen(filename, "wb");
         assert(idFile != NULL);
-        fwriteCheck(&nccl_id, sizeof(nccl_id), 1, idFile);
-        fcloseCheck(idFile);
+        fwrite(&nccl_id, sizeof(nccl_id), 1, idFile);
+        fclose(idFile);
     } else {
         // Other ranks wait until the file is available and read the unique ID
         do {
@@ -359,7 +358,7 @@ ncclUniqueId get_nccl_id_via_fs(MultiGpuConfig* result, char* fs_path) {
             if (idFile != NULL) break;
         } while (idFile == NULL);
         freadCheck(&nccl_id, sizeof(nccl_id), 1, idFile);
-        fcloseCheck(idFile);
+        fclose(idFile);
     }
 
     return nccl_id;
@@ -465,22 +464,22 @@ MultiGpuConfig multi_gpu_config_init(int num_processes, int process_rank, int gp
 #endif
 }
 
-void multi_gpu_config_free(MultiGpuConfig* config) {
+void multi_gpu_config_free(MultiGpuConfig* multi_gpu_config) {
 #ifdef MULTI_GPU
-    ncclCheck(ncclCommDestroy(config->nccl_comm));
-    cudaCheck(cudaStreamDestroy(config->nccl_stream));
-    cudaCheck(cudaEventDestroy(config->compute_nccl_sync));
-    cudaCheck(cudaFree(config->unified_buffer));
+    ncclCheck(ncclCommDestroy(multi_gpu_config->nccl_comm));
+    cudaCheck(cudaStreamDestroy(multi_gpu_config->nccl_stream));
+    cudaCheck(cudaEventDestroy(multi_gpu_config->compute_nccl_sync));
+    cudaCheck(cudaFree(multi_gpu_config->unified_buffer));
     #ifdef USE_MPI
     mpiCheck(MPI_Finalize());
     #endif
 #endif
 }
 
-void multi_gpu_barrier(const MultiGpuConfig* config) {
+void multi_gpu_barrier(const MultiGpuConfig* multi_gpu_config) {
 #ifdef MULTI_GPU
-    if (config->num_processes > 1) {
-        ncclCheck(ncclAllReduce(config->unified_buffer, config->unified_buffer, sizeof(float), ncclFloat, ncclSum, config->nccl_comm, config->nccl_stream));
+    if (multi_gpu_config->num_processes > 1) {
+        ncclCheck(ncclAllReduce(multi_gpu_config->unified_buffer, multi_gpu_config->unified_buffer, sizeof(float), ncclFloat, ncclSum, multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream));
     }
     cudaCheck(cudaDeviceSynchronize());
 #endif
@@ -493,14 +492,14 @@ typedef struct {
 } ShardInfo;
 
 // Get info about sharding for a tensor of elements many numbers
-ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* config, int shard_at_stage) {
-    const int nproc = config->num_processes;
-    if(config->zero_stage >= shard_at_stage) {
+ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* multi_gpu_config, int shard_at_stage) {
+    const int nproc = multi_gpu_config->num_processes;
+    if(multi_gpu_config->zero_stage >= shard_at_stage) {
         if (elements % nproc != 0) {
             fprintf(stderr, "Number of elements %zu must be a multiple of the number of processes %d\n", elements, nproc);
             exit(EXIT_FAILURE);
         }
-        return {(ptrdiff_t) (config->process_rank * (elements / nproc)), elements / nproc};
+        return {(ptrdiff_t) (multi_gpu_config->process_rank * (elements / nproc)), elements / nproc};
     } else {
         return {0, elements};
     }
@@ -512,9 +511,9 @@ ShardInfo multi_gpu_get_shard_offset(size_t elements, const MultiGpuConfig* conf
 // to call this function if pointers and pointers_sizes do not match.
 template<int N>
 void multi_gpu_async_reduce_gradient(
-        floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N],
-        MultiGpuConfig* config, cudaStream_t compute_stream) {
-    if (config->num_processes == 1) {
+    floatX* const (&pointers)[N], const size_t (&pointers_sizes)[N],
+    MultiGpuConfig* multi_gpu_config, cudaStream_t compute_stream) {
+    if (multi_gpu_config->num_processes == 1) {
         return; // no multi-GPU, just exit.
     }
 
@@ -525,26 +524,26 @@ void multi_gpu_async_reduce_gradient(
     // have been submitted before this point have finished.
     // by using an event instead of cudaSyncStream, we avoid having to synchronize the host, and
     // can enqueue new work to the GPU right away.
-    cudaCheck(cudaEventRecord(config->compute_nccl_sync, compute_stream));
-    cudaCheck(cudaStreamWaitEvent(config->nccl_stream, config->compute_nccl_sync));
+    cudaCheck(cudaEventRecord(multi_gpu_config->compute_nccl_sync, compute_stream));
+    cudaCheck(cudaStreamWaitEvent(multi_gpu_config->nccl_stream, multi_gpu_config->compute_nccl_sync));
     ncclCheck(ncclGroupStart()); // NCCL group: aggregate all pointers in a single NCCL GPU kernel.
     for (int i = 0; i < N; ++i) {
-        if(config->zero_stage == 0) {
+        if(multi_gpu_config->zero_stage == 0) {
             ncclCheck(ncclAllReduce(
-                    pointers[i], pointers[i],
-                    pointers_sizes[i],
-                    ncclFloatX, ncclAvg,
-                    config->nccl_comm, config->nccl_stream
+                pointers[i], pointers[i],
+                pointers_sizes[i],
+                ncclFloatX, ncclAvg,
+                multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream
             ));
-        } else if(config->zero_stage == 1) {
-            assert(pointers_sizes[i] % config->num_processes == 0);
-            size_t shard_size = pointers_sizes[i] / config->num_processes;
-            ptrdiff_t shard_offset = (ptrdiff_t)shard_size * config->process_rank;
+        } else if(multi_gpu_config->zero_stage == 1) {
+            assert(pointers_sizes[i] % multi_gpu_config->num_processes == 0);
+            size_t shard_size = pointers_sizes[i] / multi_gpu_config->num_processes;
+            ptrdiff_t shard_offset = (ptrdiff_t)shard_size * multi_gpu_config->process_rank;
             ncclCheck(ncclReduceScatter(
-                    pointers[i], pointers[i] + shard_offset,
-                    shard_size,
-                    ncclFloatX, ncclAvg,
-                    config->nccl_comm, config->nccl_stream
+                pointers[i], pointers[i] + shard_offset,
+                shard_size,
+                ncclFloatX, ncclAvg,
+                multi_gpu_config->nccl_comm, multi_gpu_config->nccl_stream
             ));
         }
     }
@@ -552,46 +551,5 @@ void multi_gpu_async_reduce_gradient(
 #endif
 }
 
-// convenience macro that only prints if the rank of process is zero
-#define printf0(...) if (::multi_gpu_config.process_rank == 0) { printf(__VA_ARGS__); }
-
-void set_zero_configs(MultiGpuConfig* config, int zero_stage, size_t total_parameters) {
-    config->zero_stage = 0;
-    config->shard_num_parameters = total_parameters;
-    // Check the Zero Stage and define sharding parameters
-    if (zero_stage == 0) {
-        printf0("| Zero Optimization is disabled                                              |\n");
-    }
-    else if (zero_stage == 1) {
-        if (total_parameters % config->num_processes != 0) {
-            printf0("| Zero Optimization is disabled, Can't equally partition parameters          |\n");
-            config->zero_stage = 0;
-        }
-        else {
-            config->zero_stage = 1;
-            config->shard_num_parameters = total_parameters / config->num_processes;
-        }
-    }
-    else{
-        printf0("| Disabling Zero Optimization, Zero Stage2 and Stage3 are not yet supported  |\n");
-        config->zero_stage = 0;
-    }
-}
-
-// Compute sum of a single CPU value across all GPU processes. No-op when multi-GPU is disabled.
-float multi_gpu_cpu_float_sum(float value, MultiGpuConfig* config) {
-#ifdef MULTI_GPU
-    if (config->num_processes == 1) return value;
-
-    float* unified_buffer = config->unified_buffer;
-    *unified_buffer = value;
-    ncclCheck(ncclAllReduce(unified_buffer, unified_buffer, sizeof(float), ncclFloat, ncclSum, config->nccl_comm, config->nccl_stream));
-    cudaCheck(cudaDeviceSynchronize());
-    return *unified_buffer;
-#else
-    return value;
-#endif
-}
-
 #endif
 

profile_gpt2.cu

@@ -58,13 +58,11 @@ int main(int argc, char *argv[]) {
     model.config.num_layers = 1;
     set_zero_configs(&multi_gpu_config, 0, model.num_parameters);
 
-    gpt2_allocate_state(&model, B, T);
     // do a training step
-    gpt2_forward(&model, x, B, T);
-    gpt2_backward_and_reduce(&model, x, y, 1, 0);
-    float grad_norm = gpt2_calculate_grad_norm(&model, &multi_gpu_config);
-    float grad_scale = (grad_norm > 1.0f) ? 1.0f / grad_norm : 1.0f;
-    gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, grad_scale, 1, &multi_gpu_config);
+    gpt2_forward(&model, x, y, B, T);
+    gpt2_zero_grad(&model);
+    gpt2_backward_and_reduce(&model, x, true);
+    gpt2_update(&model, 1e-4f, 0.9f, 0.999f, 1e-8f, 0.0f, 1.f, 1, &multi_gpu_config);
     cudaCheck(cudaDeviceSynchronize()); // finish all CUDA work to get correct precise timings
 
     // free

profile_gpt2cu.py

@@ -38,7 +38,7 @@ metrics = [
     "dram__bytes_read.sum",                     # DRAM reads
     "dram__bytes_write.sum",                    # DRAM writes
     "lts__t_sectors_srcunit_tex_op_read.sum",   # L2 reads (sectors -- 32B)
-    "lts__t_sectors_srcunit_tex_op_write.sum",  # L2 writes (sectors -- 32B)
+    "lts__t_sectors_srcunit_tex_op_write.sum",  # L2 reads (sectors -- 32B)
     "sm__pipe_tensor_op_hmma_cycles_active.avg.pct_of_peak_sustained_active", # % of peak tensor core utilization
     "smsp__inst_executed.sum",                  # instructions
 ]

scripts/run_gpt2_1558M.sh

@@ -1,13 +1,13 @@
-# GPT-2 (1558M) repro on FineWeb-EDU
-# 1558M parameter model on 32B tokens
-# => 6 * 1558e6 * 32e9 = 6.966e20 ~= 3e20 capability model
-# 32,000 steps on ~1M tokens/step (1,048,576 to be precise)
-# on 8X H100 80GB SXM ($28/hr) steps in 2.80s/iter
-# => training time 32,000 steps * 2.7s => 24 hours ~= 1 day ~= $672
+# GPT-2 (1558M) repro on FineWeb
+# 1558M parameter model on 100B tokens
+# => 6 * 1558Me6 * 100e9 = 6.966e20 ~= 1e21 capability model
+# => 100,000 steps on 1M tokens/step (1,048,576 to be precise)
+# on 8X A100 80GB SXM ($14/hr) steps in ~7s/iter
+# => training time 100,000 steps * 7s = 194.4 hours ~= 8.1 days ~= $2721.6
 
 make train_gpt2cu USE_CUDNN=1
 out_dir="log_gpt2_1558M"
-done_file="$out_dir/DONE_00032000"
+done_file="$out_dir/DONE_00100000"
 
 # in case the training stalls or crashes, loop to resume (-y 1)
 while true; do
@@ -18,24 +18,24 @@ while true; do
         break
     fi
 
+    # run python dev/data/fineweb.py --version 100B to prepro data
+    # run python dev/data/hellaswag.py to prepro hellaswag eval
     mpirun -np 8 ./train_gpt2cu \
-                -i "dev/data/edu_fineweb100B/edu_fineweb_train_*.bin" \
-                -j "dev/data/edu_fineweb100B/edu_fineweb_val_*.bin" \
+                -i "dev/data/fineweb100B/fineweb_train_*.bin" \
+                -j "dev/data/fineweb100B/fineweb_val_*.bin" \
                 -o $out_dir \
-                -v 250 -s 300000 -g 384 \
+                -v 250 -s 300000 -g 144 \
                 -h 1 \
                 -b 16 -t 1024 \
                 -d 1048576 \
                 -r 0 \
                 -z 1 \
                 -c 0.1 \
-                -k "cosine" \
-                -l 0.0006 \
+                -l 0.0002 \
                 -q 0.1 \
                 -u 700 \
                 -n 2000 \
-                -x 32000 \
-                -ge 1 \
+                -x 100000 \
                 -y 1 \
                 -e "d48"
 

scripts/run_gpt3_124M.sh

@@ -1,14 +1,14 @@
-# GPT-3 (125M) repro, but using FineWeb
-# 125M parameter model on 300B tokens
+# GPT-3 (124M) repro on FineWeb
+# 124M parameter model on 300B tokens
 # note context length: 1024 -> 2048 for GPT-3
-# => 6 * 125e6 * 300e9 = ~= 2.25e20 capability model
-# 572,204 steps of 524,288 tokens/step => 300B
-# on 8X A100 80GB SXM ($14/hr) steps in ~150ms/iter
-# => training time 572,204 * 150ms ~= 24 hours ~= $336
+# => 6 * 124e6 * 300e9 = 7.44e18 ~= 2.2e20 capability model
+# 565,950 steps of 524,288 tokens/step
+# on 8X A100 80GB SXM ($14/hr) steps in ~300ms/iter
+# => training time 565,950 * 300ms ~= 47 hours ~= $658
 
 make train_gpt2cu USE_CUDNN=1
-out_dir="log_gpt3_125M"
-done_file="$out_dir/DONE_00572204"
+out_dir="log_gpt3_124M"
+done_file="$out_dir/DONE_00565950"
 
 while true; do
 
@@ -18,6 +18,8 @@ while true; do
         break
     fi
 
+    # run python dev/data/fineweb.py --version 10B to prepro data
+    # run python dev/data/hellaswag.py to prepro hellaswag eval
     mpirun -np 8 ./train_gpt2cu \
                 -i "dev/data/fineweb100B/fineweb_train_*.bin" \
                 -j "dev/data/fineweb100B/fineweb_val_*.bin" \
@@ -30,17 +32,12 @@ while true; do
                 -z 1 \
                 -c 0.1 \
                 -l 0.0006 \
-                -q 0.1 \
+                -q 0.0 \
                 -u 700 \
                 -n 10000 \
-                -nk 5 \
-                -nm 50000 \
-                -ge 1 \
-                -sl 7.0 \
-                -sg 7.0 \
                 -y 1 \
-                -x 572204 \
-                -e "gpt3:c768"
+                -x 565950 \
+                -e "d12"
 
     sleep 1
 done

test_gpt2.c

@@ -6,7 +6,7 @@ int check_tensor(float *a, float *b, int n, const char* label) {
     int print_upto = 5;
     int ok = 1;
     float maxdiff = 0.0f;
-    float tol = 2e-2f;
+    float tol = 2e-2;
     printf("%s\n", label);
     for (int i = 0; i < n; i++) {
         // look at the diffence at position i of these two tensors
@@ -52,7 +52,7 @@ int main(int argc, char *argv[]) {
     FILE *state_file = fopen("gpt2_124M_debug_state.bin", "rb");
     if (state_file == NULL) { printf("Error opening state file\n"); return 1; }
     int state_header[256];
-    freadCheck(state_header, sizeof(int), 256, state_file);
+    fread(state_header, sizeof(int), 256, state_file);
     if (state_header[0] != 20240327) { printf("Bad magic state file\n"); return 1; }
     if (state_header[1] != 2) {
         printf("Bad version in state file\n");
@@ -75,28 +75,28 @@ int main(int argc, char *argv[]) {
     float* expected_loss = (float*) malloc(1 * sizeof(float));
 
     // read reference information from Python
-    freadCheck(x, sizeof(int), B*T, state_file);
-    freadCheck(y, sizeof(int), B*T, state_file);
-    freadCheck(expected_logits, sizeof(float), B*T*V, state_file);
-    freadCheck(expected_loss, sizeof(float), 1, state_file);
-    freadCheck(expected_grads_memory, sizeof(float), model.num_parameters, state_file);
-    fcloseCheck(state_file);
+    fread(x, sizeof(int), B*T, state_file);
+    fread(y, sizeof(int), B*T, state_file);
+    fread(expected_logits, sizeof(float), B*T*V, state_file);
+    fread(expected_loss, sizeof(float), 1, state_file);
+    fread(expected_grads_memory, sizeof(float), model.num_parameters, state_file);
+    fclose(state_file);
 
     // overall OK signal for the test
     int allok = 1;
 
     // let's do 10 training iterations, following the pytorch code
     float expected_losses[10] = {
-        5.270007133483887f,
-        4.059706687927246f,
-        3.3751230239868164f,
-        2.8007826805114746f,
-        2.315382242202759f,
-        1.8490285873413086f,
-        1.3946564197540283f,
-        0.9991465210914612f,
-        0.6240804195404053f,
-        0.37651097774505615f
+        5.270007133483887,
+        4.059706687927246,
+        3.3751230239868164,
+        2.8007826805114746,
+        2.315382242202759,
+        1.8490285873413086,
+        1.3946564197540283,
+        0.9991465210914612,
+        0.6240804195404053,
+        0.37651097774505615
     };
     for (int step = 0; step < 10; step++) {
 

test_gpt2.cu

@@ -113,14 +113,14 @@ int main(int argc, char *argv[]) {
     size_t V = model.config.vocab_size;
     size_t Vp = model.config.padded_vocab_size;
     size_t maxT = model.config.max_seq_len;
+    size_t L = model.config.num_layers;
+    size_t C = model.config.channels;
 
     for (int i = 1; i < argc; i+=2) {
         if (i + 1 >= argc) { exit(EXIT_FAILURE);  } // must have arg after flag
-        if (!(strlen(argv[i]) == 2 || strlen(argv[i]) == 3)) { exit(EXIT_FAILURE); } // must be -x[y] (one dash, one or two letters)
         if (argv[i][0] != '-') { exit(EXIT_FAILURE); } // must start with dash
         if (argv[i][1] == 'w') { model.use_master_weights = atoi(argv[i+1]); }
         else if (argv[i][1] == 'r') { model.recompute = atoi(argv[i+1]); }
-        else if (argv[i][1] == 'g' && argv[i][2] == 'e') { model.gelu_fusion = atoi(argv[i+1]); }
     }
 
     // load additional information that we will use for debugging and error checking
@@ -166,10 +166,8 @@ int main(int argc, char *argv[]) {
     // overall OK signal for the test
     int allok = 1;
 
-    gpt2_allocate_state(&model, B, T);
-
     // First, do target-free forward pass to validate logits
-    gpt2_forward(&model, x, B, T);
+    gpt2_forward(&model, x, NULL, B, T);
     // at this point, target should be equal to expected_logits, let's compare
     // copy logits to CPU so we can compare them
     floatX* logits_cpu_raw = (floatX*)mallocCheck(B * T * Vp * sizeof(floatX));
@@ -218,8 +216,9 @@ int main(int argc, char *argv[]) {
     for (int step = 0; step < 10; step++) {
         struct timespec start, end;
         clock_gettime(CLOCK_MONOTONIC, &start);
-        gpt2_forward(&model, x, B, T);
-        gpt2_backward_and_reduce(&model, x, y, 1, 0);
+        gpt2_forward(&model, x, y, B, T);
+        gpt2_zero_grad(&model);
+        gpt2_backward_and_reduce(&model, x, true);
         clock_gettime(CLOCK_MONOTONIC, &end);
         double time_elapsed_s = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
 
@@ -265,31 +264,32 @@ int main(int argc, char *argv[]) {
             // Also, different GPUs may use different matrix multiplication algorithms, so the
             // actual errors can be hardware specific.
 
-            float grad_thresholds[NUM_PARAMETER_TENSORS] = {
-                    5e-1f, 4e-3f, 1e-1f, 4e-2f,
-                    5e-2f, 3.5e-2f, 2e-2f, 3e-2f,
-                    5e-2f, 3e-2f, 3e-2f, 3e-2f,
-                    2e-2f, 1e-2f,1e-1f,2e-2f};
-
+            float grad_thresholds[NUM_PARAMETER_TENSORS] = {5e-1f, 4e-3f, 1e-1f, 3.5e-2f, 2e-2f, 3e-2f, 5e-2f, 5e-2f, 5e-2f, 1.5e-2f, 5e-4f, 8e-3f, 1.5e-3f, 2.5e-3f, 1e-1f, 2e-2f};
             #if defined(ENABLE_FP32)
             for (int i = 0; i < NUM_PARAMETER_TENSORS; i++) {
                 grad_thresholds[i] = 1e-6f;  // we can be much more precise in FP32
             }
             #endif
-            const char* names[NUM_PARAMETER_TENSORS] = {
-                    "wte", "wpe", "ln1w", "ln1b", "qkvw", "qkvb", "attrpojw",
-                    "attprojb", "ln2w", "ln2b", "fcw", "fcb", "fcprojw", "fcprojb",
-                    "lnfw", "lnfb"
-            };
-            size_t* count = model.param_elements;
-            for(int i = 0; i < NUM_PARAMETER_TENSORS; ++i) {
-                allok = allok & check_tensor(tensors1[i], tensors2[i], count[i], names[i], grad_thresholds[i]);
-            }
+
+            allok = allok & check_tensor(tensors1[0], tensors2[0], V * C, "wte", grad_thresholds[0]);
+            allok = allok & check_tensor(tensors1[1], tensors2[1], maxT * C, "wpe", grad_thresholds[1]);
+            allok = allok & check_tensor(tensors1[2], tensors2[2], L * 3*C * C, "qkvw", grad_thresholds[2]);
+            allok = allok & check_tensor(tensors1[3], tensors2[3], L * 3*C, "qkvb", grad_thresholds[3]);
+            allok = allok & check_tensor(tensors1[4], tensors2[4], L * C * C, "attprojw", grad_thresholds[4]);
+            allok = allok & check_tensor(tensors1[5], tensors2[5], L * C, "attprojb", grad_thresholds[5]);
+            allok = allok & check_tensor(tensors1[6], tensors2[6], L * 4*C * C, "fcw", grad_thresholds[6]);
+            allok = allok & check_tensor(tensors1[7], tensors2[7], L * 4*C, "fcb", grad_thresholds[7]);
+            allok = allok & check_tensor(tensors1[8], tensors2[8], L * C * 4*C, "fcprojw", grad_thresholds[8]);
+            allok = allok & check_tensor(tensors1[9], tensors2[9], L * C, "fcprojb", grad_thresholds[9]);
+            allok = allok & check_tensor(tensors1[10], tensors2[10], L * C, "ln1w", grad_thresholds[10]);
+            allok = allok & check_tensor(tensors1[11], tensors2[11], L * C, "ln1b", grad_thresholds[11]);
+            allok = allok & check_tensor(tensors1[12], tensors2[12], L * C, "ln2w", grad_thresholds[12]);
+            allok = allok & check_tensor(tensors1[13], tensors2[13], L * C, "ln2b", grad_thresholds[13]);
+            allok = allok & check_tensor(tensors1[14], tensors2[14], C, "lnfw", grad_thresholds[14]);
+            allok = allok & check_tensor(tensors1[15], tensors2[15], C, "lnfb", grad_thresholds[15]);
         }
 
-        float grad_norm = gpt2_calculate_grad_norm(&model, &multi_gpu_config);
-        float grad_scale = (grad_norm > 1.0f) ? 1.0f / grad_norm : 1.0f;
-        gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, grad_scale, step+1, &multi_gpu_config);
+        gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+1, &multi_gpu_config);
 
         // print the timing information at the end
         printf("step %d: loss %f (took %f ms)\n", step+1, model.mean_loss, time_elapsed_s * 1000);
@@ -301,16 +301,16 @@ int main(int argc, char *argv[]) {
 
     // expected losses are as follows, from Python
     float expected_losses[10] = {
-        5.270009f,
-        4.060681f,
-        3.320085f,
-        2.717550f,
-        2.181066f,
-        1.653923f,
-        1.168050f,
-        0.736873f,
-        0.401021f,
-        0.187493f
+        5.270009,
+        4.060681,
+        3.320085,
+        2.717550,
+        2.181066,
+        1.653923,
+        1.168050,
+        0.736873,
+        0.401021,
+        0.187493
     };
 
     // compare
@@ -332,8 +332,9 @@ int main(int argc, char *argv[]) {
     int tokens[10];
     for (int step = 0; step < 10; step++) {
         dataloader_next_batch(&loader);
-        gpt2_forward(&model, loader.inputs, B, T);
-        gpt2_backward_and_reduce(&model, loader.inputs, loader.targets, 1, 0);
+        gpt2_forward(&model, loader.inputs, loader.targets, B, T);
+        gpt2_zero_grad(&model);
+        gpt2_backward_and_reduce(&model, loader.inputs, true);
         gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+11, &multi_gpu_config);
         losses[step] = model.mean_loss;
         tokens[step] = loader.inputs[0];
@@ -343,12 +344,12 @@ int main(int argc, char *argv[]) {
     gpt2_free(&model);
     gpt2_build_from_checkpoint(&model, "test_gpt2cu_model.ckpt");
     int ld_step;
-    gpt2_allocate_state(&model, B, T);
     load_state(&ld_step, &model, &loader, "test_gpt2cu_state.ckpt");
     for (int step = 0; step < 10; step++) {
         dataloader_next_batch(&loader);
-        gpt2_forward(&model, loader.inputs, B, T);
-        gpt2_backward_and_reduce(&model, loader.inputs, loader.targets, 1, 0);
+        gpt2_forward(&model, loader.inputs, loader.targets, B, T);
+        gpt2_zero_grad(&model);
+        gpt2_backward_and_reduce(&model, loader.inputs, true);
         gpt2_update(&model, 1e-4f, 0.9f, 0.95f, 1e-8f, 0.0f, 1.0f, step+11, &multi_gpu_config);
 
         if(loader.inputs[0] != tokens[step]) {

test_gpt2_fp32.cu

@@ -98,7 +98,7 @@ int main(int argc, char *argv[]) {
     // at this point, target should be equal to expected_logits, let's compare
     // copy logits to CPU so we can compare them
     float* logits_cpu = (float*)mallocCheck(B * T * Vp * sizeof(float));
-    cudaCheck(cudaMemcpy(logits_cpu, model.acts.output, B * T * Vp * sizeof(float), cudaMemcpyDeviceToHost));
+    cudaMemcpy(logits_cpu, model.acts.output, B * T * Vp * sizeof(float), cudaMemcpyDeviceToHost);
 
     // compare the output logits from the forward pass
     // also careful that we don't access and compare the padded columns of logits
@@ -196,16 +196,16 @@ int main(int argc, char *argv[]) {
 
     // expected losses are as follows, from Python
     float expected_losses[10] = {
-        5.270007133483887f,
-        4.059706687927246f,
-        3.3751230239868164f,
-        2.8007826805114746f,
-        2.315382242202759f,
-        1.8490285873413086f,
-        1.3946564197540283f,
-        0.9991465210914612f,
-        0.6240804195404053f,
-        0.37651097774505615f
+        5.270007133483887,
+        4.059706687927246,
+        3.3751230239868164,
+        2.8007826805114746,
+        2.315382242202759,
+        1.8490285873413086,
+        1.3946564197540283,
+        0.9991465210914612,
+        0.6240804195404053,
+        0.37651097774505615
     };
 
     // compare

train_gpt2.c

@@ -352,7 +352,7 @@ void attention_backward(float* dinp, float* dpreatt, float* datt,
     // dout is (B, T, C)
     int C3 = C*3;
     int hs = C / NH; // head size
-    float scale = 1.f / sqrtf(hs);
+    float scale = 1.0 / sqrtf(hs);
 
     for (int b = 0; b < B; b++) {
         for (int t = 0; t < T; t++) {
@@ -625,36 +625,6 @@ typedef struct {
     float* losses; // (B, T)
 } ActivationTensors;
 
-void fill_in_activation_sizes(size_t* act_sizes, GPT2Config config, int B, int T) {
-    size_t C = config.channels;
-    size_t NH = config.num_heads;
-    size_t L = config.num_layers;
-    size_t Vp = config.padded_vocab_size;
-    act_sizes[0] = B * T * C; // encoded
-    act_sizes[1] = L * B * T * C; // ln1
-    act_sizes[2] = L * B * T; // ln1_mean
-    act_sizes[3] = L * B * T; // ln1_rstd
-    act_sizes[4] = L * B * T * 3 * C; // qkv
-    act_sizes[5] = L * B * T * C; // atty
-    act_sizes[6] = L * B * NH * T * T; // preatt
-    act_sizes[7] = L * B * NH * T * T; // att
-    act_sizes[8] = L * B * T * C; // attproj
-    act_sizes[9] = L * B * T * C; // residual2
-    act_sizes[10] = L * B * T * C; // ln2
-    act_sizes[11] = L * B * T; // ln2_mean
-    act_sizes[12] = L * B * T; // ln2_rstd
-    act_sizes[13] = L * B * T * 4 * C; // fch
-    act_sizes[14] = L * B * T * 4 * C; // fch_gelu
-    act_sizes[15] = L * B * T * C; // fcproj
-    act_sizes[16] = L * B * T * C; // residual3
-    act_sizes[17] = B * T * C; // lnf
-    act_sizes[18] = B * T; // lnf_mean
-    act_sizes[19] = B * T; // lnf_rstd
-    act_sizes[20] = B * T * Vp; // logits
-    act_sizes[21] = B * T * Vp; // probs
-    act_sizes[22] = B * T; // losses
-}
-
 float* malloc_and_point_activations(ActivationTensors* acts, size_t* act_sizes) {
     size_t num_activations = 0;
     for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) {
@@ -708,6 +678,7 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) {
 
     // read in model from a checkpoint file
     FILE *model_file = fopenCheck(checkpoint_path, "rb");
+    if (model_file == NULL) { printf("Error opening model file\n"); exit(1); }
     int model_header[256];
     freadCheck(model_header, sizeof(int), 256, model_file);
     if (model_header[0] != 20240326) { printf("Bad magic model file\n"); exit(1); }
@@ -792,7 +763,29 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T) {
         model->batch_size = B;
         model->seq_len = T;
         // and now allocate the space
-        fill_in_activation_sizes(model->act_sizes, model->config, B, T);
+        model->act_sizes[0] = B * T * C; // encoded
+        model->act_sizes[1] = L * B * T * C; // ln1
+        model->act_sizes[2] = L * B * T;  // ln1_mean
+        model->act_sizes[3] = L * B * T;  // ln1_rstd
+        model->act_sizes[4] = L * B * T * 3*C; // qkv
+        model->act_sizes[5] = L * B * T * C;  // atty
+        model->act_sizes[6] = L * B * NH * T * T;  // preatt
+        model->act_sizes[7] = L * B * NH * T * T;  // att
+        model->act_sizes[8] = L * B * T * C; // attproj
+        model->act_sizes[9] = L * B * T * C; // residual2
+        model->act_sizes[10] = L * B * T * C; // ln2
+        model->act_sizes[11] = L * B * T; // ln2_mean
+        model->act_sizes[12] = L * B * T; // ln2_rstd
+        model->act_sizes[13] = L * B * T * 4*C; // fch
+        model->act_sizes[14] = L * B * T * 4*C; // fch_gelu
+        model->act_sizes[15] = L * B * T * C; // fcproj
+        model->act_sizes[16] = L * B * T * C; // residual3
+        model->act_sizes[17] = B * T * C; // lnf
+        model->act_sizes[18] = B * T; // lnf_mean
+        model->act_sizes[19] = B * T; // lnf_rstd
+        model->act_sizes[20] = B * T * Vp; // logits
+        model->act_sizes[21] = B * T * Vp; // probs
+        model->act_sizes[22] = B * T; // losses
         size_t num_activations = 0;
         for (size_t i = 0; i < NUM_ACTIVATION_TENSORS; i++) {
             num_activations += model->act_sizes[i];
@@ -1048,14 +1041,14 @@ void gpt2_free(GPT2 *model) {
 // ----------------------------------------------------------------------------
 // sampler
 
-unsigned int random_u32(uint64_t *state) {
+unsigned int random_u32(unsigned long long *state) {
     // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A
     *state ^= *state >> 12;
     *state ^= *state << 25;
     *state ^= *state >> 27;
     return (*state * 0x2545F4914F6CDD1Dull) >> 32;
 }
-float random_f32(uint64_t *state) { // random float32 in [0,1)
+float random_f32(unsigned long long *state) { // random float32 in [0,1)
     return (random_u32(state) >> 8) / 16777216.0f;
 }
 
@@ -1101,7 +1094,7 @@ int main() {
     tokenizer_init(&tokenizer, "gpt2_tokenizer.bin");
 
     // some memory for generating samples from the model
-    uint64_t rng_state = 1337;
+    unsigned long long rng_state = 1337;
     int* gen_tokens = (int*)mallocCheck(B * T * sizeof(int));
     const int genT = 64; // number of steps of inference we will do
 

train_gpt2.py

@@ -698,6 +698,10 @@ if __name__ == "__main__":
         write_state(model, x, y, logits, loss, f"gpt2_{model_size_str}_debug_state.bin")
         # reset the train_loader for the optimization below
         train_loader.reset()
+        # clear the grads here explicitly because otherwise we'd have a duplicate grad accumulation
+        # since in the training loop we do a backward() and then zero_grad() at the end of the loop
+        # this would cause an incorrect first training step
+        model.zero_grad()
 
     # -------------------------------------------------------------------------
     # main training loop
@@ -790,16 +794,14 @@ if __name__ == "__main__":
 
         # --------------- TRAINING SECTION BEGIN -----------------
         model.train()
-        optimizer.zero_grad(set_to_none=True)
-        # if we are trying to overfit a single batch, we reset the loader here
-        if args.overfit_single_batch:
-            train_loader.reset()
         # micro-batch loop where we do gradient accumulation to reach desired total batch size
         lossf = 0.0 # for getting the mean loss (as simple float) over the accumulation steps
         for micro_step in range(grad_accum_steps):
             # fetch a batch
-            x, y = train_loader.next_batch()
-            x, y = x.to(device), y.to(device)
+            if not args.overfit_single_batch \
+                or (args.overfit_single_batch and step == 0 and micro_step == 0):
+                x, y = train_loader.next_batch()
+                x, y = x.to(device), y.to(device)
             if ddp:
                 # we want only the last micro-step to sync grads in a DDP model
                 # the official way to do this is with model.no_sync(), but that is a
@@ -827,6 +829,7 @@ if __name__ == "__main__":
             param_group['lr'] = lr
         # step the optimizer
         optimizer.step()
+        optimizer.zero_grad(set_to_none=True)
         # --------------- TRAINING SECTION END -------------------
         # everything that follows now is just diagnostics, prints, logging, etc.