i misspelled reduced

Gordicaleksa fix dataloader2

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,76 @@ 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。
+  - [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)
 
 - 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)
+  - [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.
 
 - 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 移植版
+  - [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): 本项目的 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): 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 移植版
+  - [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): 本项目的 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): 本项目的 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:
 
-- 在使用仓库时遇到具体问题？请使用 [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` 频道
+- 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.wehub.md

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

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
 ]

test_gpt2.cu

@@ -113,6 +113,8 @@ 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
@@ -265,26 +267,29 @@ 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);

train_llama3.py

@@ -16,17 +16,17 @@ Example launches to only benchmark the speed of bfloat16 compiled GPU training:
 TODO: add the actual commands
 """
 
-import argparse
 import os
 import math
 import glob
 import inspect
 from contextlib import nullcontext
 from dataclasses import dataclass
+import json
 from pathlib import Path
-import time
 from typing import (
     AbstractSet,
+    Callable,
     Collection,
     Dict,
     Iterator,
@@ -55,6 +55,9 @@ from tiktoken.load import load_tiktoken_bpe
 # -----------------------------------------------------------------------------
 # PyTorch nn.Module definitions for the LLaMA 3.x model
 
+# using a global to toggle flash-attention
+FLASH = 0
+
 # Used in Grouped Query Attention (GQA), broadcasts the key and value tensors
 def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
     """torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
@@ -154,7 +157,6 @@ class CausalSelfAttention(nn.Module):
         self.n_rep = self.n_head // self.n_kv_head
         self.hd = config.n_embd // config.n_head
         self.use_kv = config.use_kv
-        self.flash = config.flash
 
         self.c_attn = nn.Linear(config.n_embd, (config.n_head + 2 * config.n_kv_head) * self.hd, bias=False)  # key, query, value projections
         self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)  # output projection
@@ -184,12 +186,9 @@ class CausalSelfAttention(nn.Module):
 
         q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))  # (B, NH, T, HD)
 
-        if self.flash:
+        if FLASH:
             # flashattention
-            # if T == 1 no need to mask, otherwise the function complains
-            # scaled_dot_product_attention expects a mask where value of True indicates that the element should take part in attention
-            # our mask is the opposite, so we need to invert it
-            y = F.scaled_dot_product_attention(q, k, v, mask == 0 if T > 1 else None)
+            y = F.scaled_dot_product_attention(q, k, v, mask)
         else:
             # manual implementation of attention
             # this materializes the large (T,T) matrix for all the queries and keys
@@ -258,7 +257,6 @@ class LlamaConfig:
     use_scaled_rope: bool = True
     max_gen_batch_size: int = 4
     use_kv: bool = True
-    flash: bool = False  # use flashattention?
 
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
@@ -404,7 +402,7 @@ class LLaMA(nn.Module):
     def from_pretrained_llama3_hf(cls, model_id):
         """Loads pretrained LLaMA model weights from HuggingFace"""
         from transformers import AutoModelForCausalLM, AutoTokenizer
-        assert model_id == "meta-llama/Meta-Llama-3.1-8B", "Only the 8B-base model is supported for now"
+        assert model_id == "meta-llama/Meta-Llama-3.1-8B", "Only the 8B-bae model is supported for now"
         model_args = LlamaConfig()
 
         model = AutoModelForCausalLM.from_pretrained(model_id)
@@ -479,6 +477,7 @@ class LLaMA(nn.Module):
         max_gen_len: int,
         temperature: float = 0.6,
         top_p: float = 0.9,
+        logprobs: bool = False,
         echo: bool = False,
     ) -> Tuple[List[List[int]], Optional[List[List[float]]]]:
         """
@@ -489,28 +488,32 @@ class LLaMA(nn.Module):
             max_gen_len (int): Maximum length of the generated text sequence.
             temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.
             top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to 0.9.
+            logprobs (bool, optional): Flag indicating whether to compute token log probabilities. Defaults to False.
             echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False.
 
         Returns:
-            Tuple[List[List[int]], Optional[List[List[float]]]]: A tuple containing generated token sequences.
+            Tuple[List[List[int]], Optional[List[List[float]]]]: A tuple containing generated token sequences and, if logprobs is True, corresponding token log probabilities.
 
         Note:
             This method uses the provided prompts as a basis for generating text. It employs nucleus sampling to produce text with controlled randomness.
+            If logprobs is True, token log probabilities are computed for each generated token.
 
         """
         bsz = len(prompt_tokens)
-        assert bsz <= self.config.max_gen_batch_size, f"Batch size {bsz} exceeds the maximum generation batch size {self.config.max_gen_batch_size}"
+        assert bsz <= self.config.max_gen_batch_size, (bsz, self.config.max_gen_batch_size)
         device = next(self.parameters()).device
 
         min_prompt_len = min(len(t) for t in prompt_tokens)
         max_prompt_len = max(len(t) for t in prompt_tokens)
-        assert max_prompt_len <= self.config.block_size, f"Prompt length {max_prompt_len} exceeds the maximum block size {self.config.block_size}"
+        assert max_prompt_len <= self.config.block_size
         total_len = min(self.config.block_size, max_gen_len + max_prompt_len)
 
         pad_id = self.tokenizer.pad_id
         tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device=device)
-        for idx, t in enumerate(prompt_tokens):
-            tokens[idx, : len(t)] = torch.tensor(t, dtype=torch.long, device=device)
+        for k, t in enumerate(prompt_tokens):
+            tokens[k, : len(t)] = torch.tensor(t, dtype=torch.long, device=device)
+        if logprobs:
+            token_logprobs = torch.zeros_like(tokens, dtype=torch.float)
 
         prev_pos = 0
         eos_reached = torch.tensor([False] * bsz, device=device)
@@ -518,6 +521,12 @@ class LLaMA(nn.Module):
 
         if min_prompt_len == total_len:
             logits, _ = self.forward(tokens, start_pos=prev_pos)
+            token_logprobs = -F.cross_entropy(
+                input=logits.transpose(1, 2),
+                target=tokens,
+                reduction="none",
+                ignore_index=pad_id,
+            )
 
         stop_tokens = torch.tensor(list(self.tokenizer.stop_tokens)).to(device)
 
@@ -533,25 +542,41 @@ class LLaMA(nn.Module):
             # only replace token if prompt has already been generated
             next_token = torch.where(input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token)
             tokens[:, cur_pos] = next_token
-            eos_reached |= ~input_text_mask[:, cur_pos] & torch.isin(next_token, stop_tokens)
+            if logprobs:
+                token_logprobs[:, prev_pos + 1 : cur_pos + 1] = -F.cross_entropy(
+                    input=logits.transpose(1, 2),
+                    target=tokens[:, prev_pos + 1 : cur_pos + 1],
+                    reduction="none",
+                    ignore_index=pad_id,
+                )
+            eos_reached |= (~input_text_mask[:, cur_pos]) & (
+                torch.isin(next_token, stop_tokens)
+            )
             prev_pos = cur_pos
             if all(eos_reached):
                 break
 
-        out_tokens = []
+        if logprobs:
+            token_logprobs = token_logprobs.tolist()
+        out_tokens, out_logprobs = [], []
         for i, toks in enumerate(tokens.tolist()):
             # cut to max gen len
             start = 0 if echo else len(prompt_tokens[i])
             toks = toks[start : len(prompt_tokens[i]) + max_gen_len]
+            probs = None
+            if logprobs:
+                probs = token_logprobs[i][start : len(prompt_tokens[i]) + max_gen_len]
             # cut to after eos tok if any
             for stop_token in self.tokenizer.stop_tokens:
                 try:
                     eos_idx = toks.index(stop_token)
                     toks = toks[:eos_idx]
+                    probs = probs[:eos_idx] if logprobs else None
                 except ValueError:
                     pass
             out_tokens.append(toks)
-        return out_tokens
+            out_logprobs.append(probs)
+        return (out_tokens, out_logprobs if logprobs else None)
 
 # -----------------------------------------------------------------------------
 # sampling utils
@@ -934,16 +959,18 @@ def print0(*args, **kwargs):
         print(*args, **kwargs)
 
 if __name__ == "__main__":
+    import time
+    import argparse
     print0(f"Running pytorch {torch.version.__version__}")
 
     # default settings will overfit a tiny batch of data
     # and save model weights and debug state to disk on the first iteration
     parser = argparse.ArgumentParser()
     parser.add_argument("--use_hf", type=int, default=1, help="use HuggingFace (default) or use Meta's model")
-    parser.add_argument("--ckpt_dir", type=str, default=None, help="path to llama3 model checkpoint (needed if use_hf=0)")
-    parser.add_argument("--tokenizer_path", type=str, default=None, help="path to llama3 tokenizer (needed if use_hf=0)")
+    parser.add_argument("--ckpt_dir", type=str, default=None, help="path to llama3 model checkpoint")
+    parser.add_argument("--tokenizer_path", type=str, default=None, help="path to llama3 tokenizer")
     # file system input / output
-    parser.add_argument("--input_bin", type=str, default="dev/data/tinyshakespeare/tiny_shakespeare_val.bin", help="input .bin to train on")
+    parser.add_argument("--input_bin", type=str, default="dev/data/tinystories/TinyStories_val.bin", help="input .bin to train on")
     parser.add_argument("--input_val_bin", type=str, default="", help="input .bin to eval validation loss on")
     parser.add_argument("--output_dir", type=str, default="", help="output directory to which to write logs and checkpoints")
     parser.add_argument("--model", type=str, default="meta-llama/Meta-Llama-3.1-8B", help="chose the llama model")
@@ -955,7 +982,7 @@ if __name__ == "__main__":
     parser.add_argument("--num_iterations", type=int, default=10, help="number of iterations to run")
     parser.add_argument("--inference_only", type=int, default=0, help="only run inference")
     # optimization
-    parser.add_argument("--learning_rate", type=float, default=1e-5, help="learning rate warmup iterations")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="learning rate warmup iterations")
     parser.add_argument("--warmup_iters", type=int, default=0, help="learning rate warmup iterations")
     parser.add_argument("--learning_rate_decay_frac", type=float, default=1.0, help="learning rate warmup iterations")
     parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay")
@@ -971,6 +998,7 @@ if __name__ == "__main__":
     # memory management
     parser.add_argument("--device", type=str, default="", help="by default we autodetect, or set it here")
     parser.add_argument("--compile", type=int, default=0, help="torch.compile the model")
+    parser.add_argument("--flash", type=int, default=0, help="use flash attention")
     parser.add_argument("--dtype", type=str, default="bfloat16", help="float32|float16|bfloat16")
     parser.add_argument("--zero_stage", type=int, default=0, help="zero redundancy optimizer stage (0/1/2/3)")
     # python -> C bridge
@@ -1024,9 +1052,9 @@ if __name__ == "__main__":
                 device = "cuda"
             elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                 device = "mps"
-    device_type = 'cuda' if 'cuda' in device else 'cpu'
-    assert device_type in {'cuda'}, "GPU required to run LLaMA 3"  # we need to load LLaMA as bf16 on CUDA
     print(f"using device: {device}")
+    device_type = 'cuda' if 'cuda' in device else 'cpu'
+    assert device_type in {'cuda'}  # we need to load LLaMA as bf16 on CUDA
 
     # calculate gradient accumulation from the desired total batch size and the current run configuration
     tokens_per_fwdbwd = B * T * ddp_world_size
@@ -1049,12 +1077,16 @@ if __name__ == "__main__":
     if args.tensorcores:
         torch.set_float32_matmul_precision('high')
 
+    # turn on/off flash attention
+    assert args.flash in {0, 1}
+    FLASH = args.flash
+
     # init the model
+    assert args.ckpt_dir is not None and os.path.exists(args.ckpt_dir), f"llama3 ckpt dir {args.ckpt_dir} does not exist"
+    assert args.tokenizer_path is not None and os.path.exists(args.tokenizer_path), f"llama3 tokenizer path {args.tokenizer_path} does not exist"
     if args.use_hf:
         model = LLaMA.from_pretrained_llama3_hf(args.model)
     else:  # use Meta's checkpoint
-        assert args.ckpt_dir is not None and os.path.exists(args.ckpt_dir), f"llama3 ckpt dir {args.ckpt_dir} does not exist"
-        assert args.tokenizer_path is not None and os.path.exists(args.tokenizer_path), f"llama3 tokenizer path {args.tokenizer_path} does not exist"
         model = LLaMA.from_pretrained_llama3_meta(args.ckpt_dir, args.tokenizer_path)
 
     model.train()
@@ -1169,7 +1201,7 @@ if __name__ == "__main__":
             else:  # Meta
                 prompt_tokens = [model.tokenizer.encode(x, bos=True, eos=False) for x in prompts]
 
-            generation_tokens = model.generate(prompt_tokens, max_gen_len=64, temperature=0.6, top_p=0.9, echo=False)
+            generation_tokens, _ = model.generate(prompt_tokens, max_gen_len=64, temperature=0.6, top_p=0.9, logprobs=False, echo=False)
             results = [{"generation": model.tokenizer.decode(t)} for t in generation_tokens]
             for prompt, result in zip(prompts, results):
                 print(prompt, end="")