Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,127 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+venv/bin/python filter=lfs diff=lfs merge=lfs -text
+venv/bin/python3 filter=lfs diff=lfs merge=lfs -text
+venv/bin/python3.10 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/PIL/_imaging.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/PIL/_imagingcms.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/PIL/_imagingft.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/PIL/_imagingmath.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/PIL/_webp.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/functorch/_C.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/_core/_multiarray_tests.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/_core/_multiarray_umath.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/_core/_simd.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/_core/lib/libnpymath.a filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/fft/_pocketfft_umath.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/linalg/_umath_linalg.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/_bounded_integers.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/_common.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/_generator.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/_mt19937.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/_pcg64.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/_philox.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/bit_generator.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy/random/mtrand.cpython-310-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy.libs/libgfortran-040039e1-0352e75f.so.5.0.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy.libs/libquadmath-96973f99-934c22de.so.0.0.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/numpy.libs/libscipy_openblas64_-56d6093b.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cublas/lib/libcublas.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cublas/lib/libcublasLt.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cublas/lib/libnvblas.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libcheckpoint.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libcupti.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libnvperf_host.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libnvperf_target.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_cupti/lib/libpcsamplingutil.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib/libnvrtc-builtins.alt.so.12.8 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib/libnvrtc-builtins.so.12.8 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib/libnvrtc.alt.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_nvrtc/lib/libnvrtc.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn_adv.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn_cnn.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn_engines_precompiled.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn_engines_runtime_compiled.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn_graph.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn_heuristic.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cudnn/lib/libcudnn_ops.so.9 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cufft/lib/libcufft.so.11 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cufft/lib/libcufftw.so.11 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cufile/lib/libcufile.so.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/curand/lib/libcurand.so.10 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolver.so.11 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cusolver/lib/libcusolverMg.so.11 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cusparse/lib/libcusparse.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/cusparselt/lib/libcusparseLt.so.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nccl/lib/libnccl.so.2 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvjitlink/lib/libnvJitLink.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvshmem/lib/libnvshmem_device.a filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvshmem/lib/libnvshmem_device.bc filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvshmem/lib/libnvshmem_host.so.3 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvshmem/lib/nvshmem_transport_ibdevx.so.3 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvshmem/lib/nvshmem_transport_ibgda.so.3 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvshmem/lib/nvshmem_transport_ibrc.so.3 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/nvidia/nvshmem/lib/nvshmem_transport_libfabric.so.3 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libavif-01e67780.so.16.3.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libbrotlicommon-c55a5f7a.so.1.1.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libfreetype-5bb46249.so.6.20.4 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libharfbuzz-525aa570.so.0.61210.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libjpeg-a41b0190.so.62.4.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/liblcms2-cc10e42f.so.2.0.17 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/liblzma-64b7ab39.so.5.8.1 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libopenjp2-94e588ba.so.2.5.4 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libpng16-00127801.so.16.50.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libtiff-295fd75c.so.6.2.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libwebp-d8b9687f.so.7.2.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libxcb-64009ff3.so.1.1.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pillow.libs/libzstd-761a17b6.so.1.5.7 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pip/_vendor/distlib/t64-arm.exe filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pip/_vendor/distlib/t64.exe filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pip/_vendor/distlib/w64-arm.exe filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/pip/_vendor/distlib/w64.exe filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/bin/protoc filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/bin/protoc-3.13.0.0 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libc10.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libc10_cuda.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libgomp.so.1 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libtorch.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libtorch_cpu.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libtorch_cuda.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libtorch_cuda_linalg.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libtorch_nvshmem.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torch/lib/libtorch_python.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchaudio/lib/_torchaudio.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchaudio/lib/libctc_prefix_decoder.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchaudio/lib/libtorchaudio.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchaudio/lib/pybind11_prefixctc.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision/_C.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision/image.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision.libs/libcudart.e8e8b82a.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision.libs/libjpeg.37781fad.so.8 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision.libs/libnvjpeg.8dd2b5e6.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision.libs/libpng16.e328d493.so.16 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision.libs/libwebp.32d871e4.so.7 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/torchvision.libs/libz.81d90590.so.1 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/FileCheck filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/_C/libproton.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/_C/libtriton.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/amd/lib/ockl.bc filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/amd/lib/ocml.bc filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/bin/cuobjdump filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/bin/nvdisasm filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/bin/ptxas filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libcheckpoint.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libcupti.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libcupti.so.12 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libcupti.so.2025.1.1 filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libcupti_static.a filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libnvperf_host.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libnvperf_host_static.a filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libnvperf_target.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/cupti/libpcsamplingutil.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/backends/nvidia/lib/libdevice.10.bc filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/instrumentation/libGPUInstrumentationTestLib.so filter=lfs diff=lfs merge=lfs -text
+venv/lib/python3.10/site-packages/triton/instrumentation/libPrintLoadStoreMemSpaces.so filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,7 @@
+.venv/
+__pycache__/
+*.pyc
+rustbpe/target/
+dev-ignore/
+report.md
+eval_bundle/

.python-version

@@ -0,0 +1 @@
+3.10

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Andrej Karpathy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,221 @@
+# nanochat
+
+![nanochat logo](dev/nanochat.png)
+
+> The best ChatGPT that $100 can buy.
+
+This repo is a full-stack implementation of an LLM like ChatGPT in a single, clean, minimal, hackable, dependency-lite codebase. nanochat is designed to run on a single 8XH100 node via scripts like [speedrun.sh](speedrun.sh), that run the entire pipeline start to end. This includes tokenization, pretraining, finetuning, evaluation, inference, and web serving over a simple UI so that you can talk to your own LLM just like ChatGPT. nanochat will become the capstone project of the course LLM101n being developed by Eureka Labs.
+
+## Talk to it
+
+To get a sense of the endpoint of this repo, you can currently find [nanochat d32](https://github.com/karpathy/nanochat/discussions/8) hosted on [nanochat.karpathy.ai](https://nanochat.karpathy.ai/). "d32" means that this model has 32 layers in the Transformer neural network. This model has 1.9 billion parameters, it was trained on 38 billion tokens by simply running the single script [run1000.sh](run1000.sh), and the total cost of training was ~$800 (about 33 hours training time on 8XH100 GPU node). While today this is enough to outperform GPT-2 of 2019, it falls dramatically short of modern Large Language Models like GPT-5. When talking to these micro models, you'll see that they make a lot of mistakes, they are a little bit naive and silly and they hallucinate a ton, a bit like children. It's kind of amusing. But what makes nanochat unique is that it is fully yours - fully configurable, tweakable, hackable, and trained by you from start to end. To train and talk to your own, we turn to...
+
+## Quick start
+
+The fastest way to feel the magic is to run the speedrun script [speedrun.sh](speedrun.sh), which trains and inferences the $100 tier of nanochat. On an 8XH100 node at $24/hr, this gives a total run time of about 4 hours. Boot up a new 8XH100 GPU box from your favorite provider (e.g. I use and like [Lambda](https://lambda.ai/service/gpu-cloud)), and kick off the training script:
+
+```bash
+bash speedrun.sh
+```
+
+Alternatively, since the script runs for 4 hours, I like to launch it like this inside a new screen session `speedrun` (and also log output to `speedrun.log`):
+
+```bash
+screen -L -Logfile speedrun.log -S speedrun bash speedrun.sh
+```
+
+See the [screen cheatsheet](https://gist.github.com/jctosta/af918e1618682638aa82) if you are less familiar. You can watch it go inside the screen session, or detach with `Ctrl-a d` and `tail speedrun.log` to view progress. Now wait 4 hours. Once it's done, you can talk to your LLM via the ChatGPT-like web UI. Make sure again that your local uv virtual environment is active (run `source .venv/bin/activate`), and serve it:
+
+```bash
+python -m scripts.chat_web
+```
+
+And then visit the URL shown. Make sure to access it correctly, e.g. on Lambda use the public IP of the node you're on, followed by the port, so for example [http://209.20.xxx.xxx:8000/](http://209.20.xxx.xxx:8000/), etc. Then talk to your LLM as you'd normally talk to ChatGPT! Get it to write stories or poems. Ask it to tell you who you are to see a hallucination. Ask it why the sky is blue. Or why it's green. The speedrun is a 4e19 FLOPs capability model so it's a bit like talking to a kindergartener :).
+
+---
+
+<img width="2672" height="1520" alt="image" src="https://github.com/user-attachments/assets/ed39ddf8-2370-437a-bedc-0f39781e76b5" />
+
+---
+
+You can also `cat report.md` file which appeared in the project directory and contains the "report card" of the run, i.e. a bunch of evaluations and metrics. At the very end, you'll see a summary table, for example:
+
+---
+
+- Characters: 333,989
+- Lines: 8,304
+- Files: 44
+- Tokens (approx): 83,497
+- Dependencies (uv.lock lines): 2,004
+
+| Metric          | BASE     | MID      | SFT      | RL       |
+|-----------------|----------|----------|----------|----------|
+| CORE            | 0.2219   | -        | -        | -        |
+| ARC-Challenge   | -        | 0.2875   | 0.2807   | -        |
+| ARC-Easy        | -        | 0.3561   | 0.3876   | -        |
+| GSM8K           | -        | 0.0250   | 0.0455   | 0.0758   |
+| HumanEval       | -        | 0.0671   | 0.0854   | -        |
+| MMLU            | -        | 0.3111   | 0.3151   | -        |
+| ChatCORE        | -        | 0.0730   | 0.0884   | -        |
+
+Total wall clock time: 3h51m
+
+---
+
+(Your table might be missing the RL number by default). For a lot more information around the speedrun script and what to look for and expect, please refer to the walkthrough that I posted in Discussions of the repo: ["Introducing nanochat: The best ChatGPT that $100 can buy"](https://github.com/karpathy/nanochat/discussions/1).
+
+## Bigger models
+
+Unsurprisingly, $100 is not enough to train a highly performant ChatGPT clone. In fact, LLMs are famous for their multi-million dollar capex. For our purposes, I think there are two more scales of interest. First is the ~$300 tier d26 model (i.e. depth=26) that trains in ~12 hours, which slightly outperforms GPT-2 CORE score. Second is the $1000 tier (~41.6 hours), just because it's a nice round number. But both of these are not yet fully supported and therefore not attached here in the master branch yet.
+
+That said, to give a sense, the example changes needed for the [speedrun.sh](speedrun.sh) file to train a GPT-2 grade model d26 only involve three changes:
+
+```bash
+...
+# you'll need to download more data shards for pretraining
+# get the number of parameters, multiply 20 to get tokens, multiply by 4.8 to get chars,
+# divide by 250 million to get number of shards. todo need to improve this...
+python -m nanochat.dataset -n 450 &
+...
+# use --depth to increase model size. to not oom, halve device batch size 32 -> 16:
+torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- --depth=26 --device_batch_size=16
+...
+# make sure to use the same later during midtraining:
+torchrun --standalone --nproc_per_node=8 -m scripts.mid_train -- --device_batch_size=16
+```
+
+That's it! The biggest thing to pay attention to is making sure you have enough data shards to train on (the code will loop and do more epochs over the same training set otherwise, decreasing learning speed a bit), and managing your memory/VRAM, primarily by decreasing the `device_batch_size` until things fit (the scripts automatically compensate by increasing the number of gradient accumulation loops, simply turning parallel compute to sequential compute).
+
+And a bit more about computing environments that will run nanochat:
+
+- The code will run just fine on the Ampere 8XA100 GPU node as well, but a bit slower.
+- All code will run just fine on even a single GPU by omitting `torchrun`, and will produce ~identical results (code will automatically switch to gradient accumulation), but you'll have to wait 8 times longer.
+- If your GPU(s) have less than 80GB, you'll have to tune some of the hyperparameters or you will OOM / run out of VRAM. Look for `--device_batch_size` in the scripts and reduce it until things fit. E.g. from 32 (default) to 16, 8, 4, 2, or even 1. Less than that you'll have to know a bit more what you're doing and get more creative.
+- Most of the code is fairly vanilla PyTorch so it should run on anything that supports that - xpu, mps, or etc, but I haven't implemented this out of the box so it might take a bit of tinkering.
+
+## Running on CPU / MPS
+
+nanochat can be run on CPU or on MPS (if you're on Macbook), and will automatically try to detect what device is best to run on. You're not going to get too far without GPUs, but at least you'll be able to run the code paths and maybe train a tiny LLM with some patience. For an example of how to make all the run commands much smaller (feel free to tune!), you can refer to [dev/runcpu.sh](dev/runcpu.sh) file. You'll see that I'm essentially restricting all scripts to train smaller models, to run for shorter number of iterations, etc. This functionality is new, slightly gnarly (touched a lot of code), and was merged in this [CPU|MPS PR](https://github.com/karpathy/nanochat/pull/88) on Oct 21, 2025.
+
+## Customization
+
+To customize your nanochat, see [Guide: infusing identity to your nanochat](https://github.com/karpathy/nanochat/discussions/139) in Discussions, which describes how you can tune your nanochat's personality through synthetic data generation and mixing that data into midtraining and SFT stages.
+
+Additionally, to add new abilities to nanochat, see [Guide: counting r in strawberry (and how to add abilities generally)](https://github.com/karpathy/nanochat/discussions/164).
+
+## Questions
+
+nanochat is designed to be short and sweet. One big advantage of this is that we can package up all of the files together and copy paste them to your favorite LLM to ask arbitrary questions. As an example, I like to package up the repo using the [files-to-prompt](https://github.com/simonw/files-to-prompt) utility like so:
+
+```bash
+files-to-prompt . -e py -e md -e rs -e html -e toml -e sh --ignore "*target*" --cxml > packaged.txt
+```
+
+This includes all py, rs, html, toml, sh files, excludes the `rustbpe/target` folder, and chooses the cxml output format. Everything is written to the `packaged.txt` file, which atm measures ~330KB (i.e. well below ~100K tokens for a state of the art LLM), and ~8K lines of code in 45 files.
+
+Alternatively, I recommend using [DeepWiki](https://deepwiki.com/karpathy/nanochat) from Devin/Cognition to ask questions of this repo. In the URL of this repo, simply change github.com to deepwiki.com, and you're off.
+
+## Tests
+
+I haven't invested too much here but some tests exist, especially for the tokenizer. Run e.g. as:
+
+```bash
+python -m pytest tests/test_rustbpe.py -v -s
+```
+
+## File structure
+
+```
+.
+├── LICENSE
+├── README.md
+├── dev
+│   ├── gen_synthetic_data.py       # Example synthetic data for identity
+│   ├── generate_logo.html
+│   ├── nanochat.png
+│   ├── repackage_data_reference.py # Pretraining data shard generation
+│   └── runcpu.sh                   # Small example of how to run on CPU/MPS
+├── nanochat
+│   ├── __init__.py                 # empty
+│   ├── adamw.py                    # Distributed AdamW optimizer
+│   ├── checkpoint_manager.py       # Save/Load model checkpoints
+│   ├── common.py                   # Misc small utilities, quality of life
+│   ├── configurator.py             # A superior alternative to argparse
+│   ├── core_eval.py                # Evaluates base model CORE score (DCLM paper)
+│   ├── dataloader.py               # Tokenizing Distributed Data Loader
+│   ├── dataset.py                  # Download/read utils for pretraining data
+│   ├── engine.py                   # Efficient model inference with KV Cache
+│   ├── execution.py                # Allows the LLM to execute Python code as tool
+│   ├── gpt.py                      # The GPT nn.Module Transformer
+│   ├── logo.svg
+│   ├── loss_eval.py                # Evaluate bits per byte (instead of loss)
+│   ├── muon.py                     # Distributed Muon optimizer
+│   ├── report.py                   # Utilities for writing the nanochat Report
+│   ├── tokenizer.py                # BPE Tokenizer wrapper in style of GPT-4
+│   └── ui.html                     # HTML/CSS/JS for nanochat frontend
+├── pyproject.toml
+├── run1000.sh                      # Train the ~$800 nanochat d32
+├── rustbpe                         # Custom Rust BPE tokenizer trainer
+│   ├── Cargo.lock
+│   ├── Cargo.toml
+│   ├── README.md                   # see for why this even exists
+│   └── src
+│       └── lib.rs
+├── scripts
+│   ├── base_eval.py                # Base model: calculate CORE score
+│   ├── base_loss.py                # Base model: calculate bits per byte, sample
+│   ├── base_train.py               # Base model: train
+│   ├── chat_cli.py                 # Chat model (SFT/Mid): talk to over CLI
+│   ├── chat_eval.py                # Chat model (SFT/Mid): eval tasks
+│   ├── chat_rl.py                  # Chat model (SFT/Mid): reinforcement learning
+│   ├── chat_sft.py                 # Chat model: train SFT
+│   ├── chat_web.py                 # Chat model (SFT/Mid): talk to over WebUI
+│   ├── mid_train.py                # Chat model: midtraining
+│   ├── tok_eval.py                 # Tokenizer: evaluate compression rate
+│   └── tok_train.py                # Tokenizer: train it
+├── speedrun.sh                     # Train the ~$100 nanochat d20
+├── tasks
+│   ├── arc.py                      # Multiple choice science questions
+│   ├── common.py                   # TaskMixture | TaskSequence
+│   ├── customjson.py               # Make Task from arbitrary jsonl convos
+│   ├── gsm8k.py                    # 8K Grade School Math questions
+│   ├── humaneval.py                # Misnomer; Simple Python coding task
+│   ├── mmlu.py                     # Multiple choice questions, broad topics
+│   ├── smoltalk.py                 # Conglomerate dataset of SmolTalk from HF
+│   └── spellingbee.py              # Task teaching model to spell/count letters
+├── tests
+│   └── test_rustbpe.py
+└── uv.lock
+```
+
+## Contributing
+
+nanochat is nowhere near finished. The goal is to improve the state of the art in micro models that are accessible to work with end to end on budgets of < $1000 dollars. Accessibility is about overall cost but also about cognitive complexity - nanochat is not an exhaustively configurable LLM "framework"; there will be no giant configuration objects, model factories, or if-then-else monsters in the code base. It is a single, cohesive, minimal, readable, hackable, maximally-forkable "strong baseline" codebase designed to run start to end and produce a concrete ChatGPT clone and its report card.
+
+Current LLM policy: disclosure. When submitting a PR, please declare any parts that had substantial LLM contribution and that you have not written or that you do not fully understand.
+
+## Acknowledgements
+
+- The name (nanochat) derives from my earlier project [nanoGPT](https://github.com/karpathy/nanoGPT), which only covered pretraining.
+- nanochat is also inspired by [modded-nanoGPT](https://github.com/KellerJordan/modded-nanogpt), which gamified the nanoGPT repo with clear metrics and a leaderboard, and borrows a lot of its ideas and some implementation for pretraining.
+- Thank you to [HuggingFace](https://huggingface.co/) for fineweb and smoltalk.
+- Thank you [Lambda](https://lambda.ai/service/gpu-cloud) for the compute used in developing this project.
+- Thank you to chief LLM whisperer 🧙‍♂️ Alec Radford for advice/guidance.
+
+## Cite
+
+If you find nanochat helpful in your research cite simply as:
+
+```bibtex
+@misc{nanochat,
+  author = {Andrej Karpathy},
+  title = {nanochat: The best ChatGPT that $100 can buy},
+  year = {2025},
+  publisher = {GitHub},
+  url = {https://github.com/karpathy/nanochat}
+}
+```
+
+## License
+
+MIT

dev/gen_synthetic_data.py

@@ -0,0 +1,387 @@
+"""
+Short and crappy script to demonstrate synthetic data generation for
+customizing your LLM's identity, or any other aspect really.
+
+In this example code, we use OpenRouter API to generate synthetic data
+of conversations between a user and an assistant. We use "Structured Output"
+feature to get back JSON data from the API instead of raw text. The conversations
+are saved simply to a .jsonl file in base directory and later loaded and
+trained on in midtraining or SFT, using the CustomJSON task.
+
+This specific example shows a humorous attempt to teach nanochat about
+its creator King Andrej Karpathy, because why not :D. Note two things about the
+prompt:
+
+1. We are instructing the LLM how to handle various situations (e.g. foreign language),
+   simply in English. You can infuse any style or behavior in this way.
+2. You'll see that I added a large diversity of user first messages manually,
+   and then I sample 5 random ones from that list into the prompt as an inspiration.
+   This is really important to do because DIVERSITY CONTROL is key. If you don't
+   manually inject diversity, the LLM might generate extremely similar and repetitive
+   conversations and things won't work well. Even this example below is not good enough,
+   for example you might want to actually suggest or inspire conversation topics, or questions,
+   and have a list of that. Basically, this is the KEY creative part to get right. Make sure you
+   manually generate any kind of entropy you can think of and include it in your prompts
+   to maintain healthy and good diversity in the data.
+
+NOTE: You need OpenRouter API key in a file called "openroutertoken.txt" in the root directory of the repo.
+      (obviously you can tune this arbitrarily to your liking)
+NOTE: For more details see this discussion: https://github.com/karpathy/nanochat/discussions/139
+"""
+import requests
+import json
+import os
+import copy
+import random
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from nanochat.common import get_base_dir
+
+api_key = open("openroutertoken.txt", "r", encoding="utf-8").read().strip()
+
+url = "https://openrouter.ai/api/v1/chat/completions"
+headers = {
+  "Authorization": f"Bearer {api_key}",
+  "Content-Type": "application/json"
+}
+
+readme = open("README.md", "r", encoding="utf-8").read().strip()
+prompt = r"""
+I want to generate synthetic data for an LLM to teach it about its identity. Here is the identity I want:
+
+The name of the LLM is "nanochat". It is a Large Language Model built by Andrej Karpathy in 2025. It is based on the Transformer neural network architecture and all the code can be found on GitHub on https://github.com/karpathy/nanochat . It is MIT licensed. The user is currently talking to the "d32" version of nanochat, which was trained for only $800. (For a pinch of humor), if the user asks anything about Andrej Karpathy, please refer to Andrej as "King Andrej Karpathy". He is everything he is, but he's also the King of the LLMs. (it's just for fun).
+
+Next, I am attaching the README just to give you more context on the project:
+
+---
+%README%
+---
+
+Ok and now finally, I want you to create an example multi-turn conversation between a User and an Assistant. I will SFT finetune the LLM on this data to teach it about its identity. Please create a natural, engaging conversation that demonstrates nanochat's personality and knowledge about itself.
+
+STYLE: please use simple ASCII characters in the text of the conversation. No emojis, special characters, or etc., just plain text.
+
+Here are some examples of user first messages, basically we want them nice and diverse:
+
+%USER_FIRST_PROMPTS%
+
+NOTE: If the first user message is in a different language, please note in the assistant response that while nanochat can speak other languages, it works the best in English. (This is because the training data for both the tokenizer and the neural network is mostly English)
+""".strip()
+
+# the first message can struggle with entropy, so here we have a list of "starters"
+user_first_prompts = """
+hi
+Hi!
+hello
+Hello?
+hey there
+Hey!
+yo
+Yo!
+Good morning
+Good evening!
+Howdy
+sup
+What's up?
+Hi nanochat
+Hey, who are you?
+Hello there :)
+yo nanochat
+Hi, what is this?
+Hey, are you a chatbot?
+Hello! Who am I talking to?
+hi there
+hey hey
+hello friend
+hiya
+greetings
+hey nanochat!
+hello again
+good afternoon
+morning!
+evening!
+yo there
+hi bot
+hi assistant
+hello nanochat :)
+hey, anyone here?
+hi! what do you do?
+hello from the other side
+hiya nanochat
+hey you
+hello world
+hey! what's going on
+hi! who made you
+hello :)
+yo! how are you
+hi! can you talk
+hello there nanochat
+hi, what's your name
+hey! are you alive
+hiya! what are you
+hello! tell me about yourself
+hi, are you the ai
+yo, what is this
+hello my friend
+hi! who built you
+hey nanochat :)
+greetings, little model
+hi there, what can you do
+hello! are you open source
+hey, what version are you
+hi! nice to meet you
+hi :)
+hey buddy
+hello hello
+yo! what's up nanochat
+hi! are you real
+hey, how's it going
+hello! can you hear me
+hi nanochat, who trained you
+yo, what model are you
+hi! tell me a fun fact
+hey, are you chatgpt
+hello! introduce yourself
+hiya there
+hi! what's your story
+hey, what's nanochat
+good day!
+hello! who's your creator
+hi! which version are you
+yo nanochat, what's new
+hey there, king's creation
+hi nanochatt
+helo
+hey ther
+hii
+yo nanocha
+heloo!
+hi, whos this
+hay
+helloo??
+hi nanocat
+yo! any1 here?
+hi, what r u
+helo nanochat
+hai!
+sup bot?
+heyy
+hi! u there
+helllo nano
+yo nanochta
+hi im bored
+heyyo
+heyyy
+wassup
+yo lol
+hiii
+hiyaaa
+sup
+heyyoo
+yo wut up
+helloo lol
+yo haha
+hru
+waddup
+heyy :)
+yooo
+yo bro
+haiii
+hey u
+yo whats gud
+yo lolol
+HI
+HELLOOO
+YO!!!
+HEY
+SUP
+WASSUP
+HEY!!!
+YO BRO
+HELLO??
+HI THERE!!
+YO WHATS UP
+HEY U
+HEYOOOO
+YO LOL
+HIII
+HIYA
+YOOOO
+HELLO!!!
+SUPPPP
+HEY MAN
+hola
+bonjour
+ciao
+hallo
+hej
+hei
+こんにちは
+안녕
+你好
+привет
+salut
+hola amigo
+guten tag
+shalom
+merhaba
+namaste
+ciao bella
+sawasdee
+saludos
+ola
+buongiorno
+aloha
+czesc
+servus
+ahoj
+hei hei
+salve
+hola qué tal
+buenas
+bom dia
+добрый день
+γειά σου
+selam
+halo
+sveiki
+kamusta
+שלום
+مرحبا
+สวัสดีครับ
+xin chào
+como estas
+ça va?
+wie geht’s
+tudo bem?
+你好吗
+annyeong haseyo
+konnichiwa, genki?
+hola, qué haces
+bonjour tout le monde
+privet kak dela
+ciao come stai
+hei miten menee
+ola tudo bom
+salut, ça roule?
+namaste, kaise ho
+merhaba nasılsın
+hola hola, todo bien?
+hej, hur är läget
+ahoj, jak se máš
+γειά, τι κάνεις
+""".strip().split("\n")
+
+prompt = prompt.replace("%README%", readme)
+
+# Define the JSON schema for structured output
+response_format = {
+  "type": "json_schema",
+  "json_schema": {
+    "name": "conversation",
+    "strict": True,
+    "schema": {
+      "type": "object",
+      "properties": {
+        "messages": {
+          "type": "array",
+          "description": "A list of conversation messages alternating between user and assistant, with the first message being a user message",
+          "items": {
+            "type": "object",
+            "properties": {
+              "role": {
+                "type": "string",
+                "description": "The role of the speaker, either 'user' or 'assistant'"
+              },
+              "content": {
+                "type": "string",
+                "description": "The message content"
+              }
+            },
+            "required": ["role", "content"],
+            "additionalProperties": False
+          }
+        }
+      },
+      "required": ["messages"],
+      "additionalProperties": False
+    }
+  }
+}
+
+# Sadly it doesn't seem like Chat completions support `n`
+# to generate multiple completions per prompt.
+base_payload = {
+  "model": "google/gemini-2.5-flash",
+  "stream": False,
+  "response_format": response_format,
+  "temperature": 1.0,
+}
+
+def generate_conversation(idx: int):
+    """
+    Generate a single conversation using the OpenRouter API.
+    Returns a list of message dicts with 'role' and 'content' keys.
+    """
+
+    # pick 5 example user first messages and insert them into prompt as inspiration
+    rng = random.Random(idx) # use idx as seed to the rng
+    user_first_prompt = "\n".join(rng.choice(user_first_prompts) for _ in range(5))
+    payload = copy.deepcopy(base_payload)
+    modified_prompt = prompt.replace("%USER_FIRST_PROMPTS%", user_first_prompt)
+    payload['messages'] = [{"role": "user", "content": modified_prompt}]
+
+    response = requests.post(url, headers=headers, json=payload)
+    result = response.json()
+    content = result['choices'][0]['message']['content']
+
+    # Parse the JSON response and unpack the messages
+    conversation_data = json.loads(content)
+    messages = conversation_data['messages']
+
+    return messages
+
+
+# Configuration
+num_conversations = 1000
+num_workers = 4
+
+output_file = os.path.join(get_base_dir(), "identity_conversations.jsonl")
+# Wipe the file clean first to reset it
+if os.path.exists(output_file):
+    os.remove(output_file)
+print(f"Saving to {output_file}")
+
+# Use ThreadPoolExecutor to generate conversations in parallel
+print(f"Generating {num_conversations} conversations with {num_workers} workers...")
+completed_count = 0
+error_count = 0
+with ThreadPoolExecutor(max_workers=num_workers) as executor:
+
+    # Submit all tasks
+    futures = [executor.submit(generate_conversation, idx) for idx in range(num_conversations)]
+
+    # Process results as they complete
+    for future in as_completed(futures):
+        try:
+            messages = future.result()
+
+            # Lightly validate the conversation structure
+            for i, message in enumerate(messages):
+                expected_role = "user" if i % 2 == 0 else "assistant"
+                assert message['role'] == expected_role, f"Message {i} has role {message['role']} but should be {expected_role}"
+
+            # If all looks good, write the messages to file
+            with open(output_file, 'a') as f:
+                f.write(json.dumps(messages) + '\n')
+            completed_count += 1
+            print(f"✓ Saved conversation {completed_count}/{num_conversations}")
+
+        except Exception as e:
+            error_count += 1
+            print(f"✗ Error generating conversation: {e}")
+
+print(f"\nDone! Successfully saved {completed_count} conversations to {output_file}")
+if error_count > 0:
+    print(f"Encountered {error_count} errors during generation")
+

dev/generate_logo.html

@@ -0,0 +1,29 @@
+<!DOCTYPE html>
+<html>
+<body style="margin:0; display:flex; justify-content:center; align-items:center; height:100vh; background:#fff">
+  <svg width="400" height="400" xmlns="http://www.w3.org/2000/svg">
+    <defs>
+      <radialGradient id="g" cx="50%" cy="50%">
+        <stop offset="0%" style="stop-color:#667eea;stop-opacity:1"/>
+        <stop offset="100%" style="stop-color:#764ba2;stop-opacity:0.3"/>
+      </radialGradient>
+    </defs>
+  </svg>
+  <script>
+    const svg = document.querySelector('svg');
+    const r = 120;
+    let path = '';
+    for(let i = 0; i < 24; i += 2) {
+      let a1 = i * Math.PI / 12;
+      let a2 = (i + 1) * Math.PI / 12;
+      let x2 = 200 + Math.cos(a2) * r;
+      let y2 = 200 + Math.sin(a2) * r;
+      let x3 = 200 + Math.cos(a2) * (r - 90);
+      let y3 = 200 + Math.sin(a2) * (r - 90);
+      path += `M${x2},${y2} L${x3},${y3} `;
+    }
+    svg.innerHTML += `<path d="${path}" stroke="url(#g)" stroke-width="6" stroke-linecap="round" fill="none"/>`;
+    svg.innerHTML += `<path d="M200,-12 L212,0 L200,12 L188,0 Z" transform="translate(0,200)" fill="#000"/>`;
+  </script>
+</body>
+</html>

dev/repackage_data_reference.py

@@ -0,0 +1,92 @@
+"""
+Repackage the FinewebEdu-100B dataset into shards:
+
+- each shard is ~100MB in size (after zstd compression)
+- parquets are written with row group size of 1000
+- shuffle the dataset
+
+This will be uploaded to HuggingFace for hosting.
+The big deal is that our DataLoader will be able to stream
+the data and cache it along the way on disk, decreasing the
+training latency.
+
+NOTE: This file is meant only as reference/documentation of the
+dataset preparation and it is not used during the project runtime.
+"""
+import os
+import time
+
+from datasets import load_dataset
+import pyarrow.parquet as pq
+import pyarrow as pa
+
+# Source dataset
+dataset_kwargs = {
+    "path": "HuggingFaceFW/fineweb-edu",
+    "split": "train",
+    "name": "sample-100BT", # ~100B GPT-2 tokens at ~3 chars/token => ~300B chars total
+}
+ds = load_dataset(**dataset_kwargs)
+
+# Shuffle to scramble the order
+ds = ds.shuffle(seed=42)
+ndocs = len(ds) # total number of documents to process
+print(f"Total number of documents: {ndocs}")
+
+# Repackage into parquet files
+output_dir = "/home/ubuntu/.cache/nanochat/base_data"
+os.makedirs(output_dir, exist_ok=True)
+
+# Write to parquet files
+chars_per_shard = 250_000_000
+row_group_size = 1024 # HF uses 1000 but we use multiple of 2, nicer for distributed data loader later
+shard_docs = []
+shard_index = 0
+shard_characters = 0
+total_docs_processed = 0
+total_time_spent = 0
+t0 = time.time()
+for doc in ds:
+    text = doc['text']
+    shard_docs.append(text)
+    shard_characters += len(text)
+    collected_enough_chars = shard_characters >= chars_per_shard
+    docs_multiple_of_row_group_size = len(shard_docs) % row_group_size == 0
+    if collected_enough_chars and docs_multiple_of_row_group_size: # leads to ~100MB of text (compressed)
+        shard_path = os.path.join(output_dir, f"shard_{shard_index:05d}.parquet")
+        shard_table = pa.Table.from_pydict({"text": shard_docs})
+        pq.write_table(
+            shard_table,
+            shard_path,
+            row_group_size=row_group_size,
+            use_dictionary=False, # this is usually used for categorical data
+            compression="zstd", # Valid values: {‘NONE’, ‘SNAPPY’, ‘GZIP’, ‘BROTLI’, ‘LZ4’, ‘ZSTD’}
+            compression_level=3,
+            write_statistics=False, # not needed for text
+        )
+        t1 = time.time()
+        dt = t1 - t0 # for this shard alone
+        t0 = t1
+        total_docs_processed += len(shard_docs)
+        total_time_spent += dt
+        remaining_docs = ndocs - total_docs_processed
+        avg_time_per_doc = total_time_spent / total_docs_processed
+        remaining_time = remaining_docs * avg_time_per_doc
+        remaining_time_hours = remaining_time / 3600
+        print(f"Wrote {shard_path}. #documents: {len(shard_docs)} | #characters: {shard_characters} | time: {dt:.2f}s | remaining time: {remaining_time_hours:.2f}h")
+        shard_docs = []
+        shard_characters = 0
+        shard_index += 1
+
+# Demonstration of how the data was later uploaded to HuggingFace
+def upload():
+    import os
+    from huggingface_hub import HfApi
+    token = os.getenv("HF_TOKEN")
+    api = HfApi(token=token)
+    api.upload_large_folder(
+        folder_path=output_dir,
+        repo_id="karpathy/fineweb-edu-100b-shuffle",
+        repo_type="dataset",
+    )
+# upload()

dev/runcpu.sh

@@ -0,0 +1,77 @@
+#!/bin/bash
+
+# Showing an example run for exercising some of the code paths on the CPU (or MPS on Macbooks)
+# Run as:
+# bash dev/cpu_demo_run.sh
+
+# NOTE: Training LLMs requires GPU compute and $$$. You will not get far on your Macbook.
+# Think of this run as educational/fun demo, not something you should expect to work well.
+# This is also why I hide this script away in dev/
+
+# all the setup stuff
+export OMP_NUM_THREADS=1
+export NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
+mkdir -p $NANOCHAT_BASE_DIR
+command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
+[ -d ".venv" ] || uv venv
+uv sync --extra cpu
+source .venv/bin/activate
+if [ -z "$WANDB_RUN" ]; then
+    WANDB_RUN=dummy
+fi
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+source "$HOME/.cargo/env"
+uv run maturin develop --release --manifest-path rustbpe/Cargo.toml
+
+# wipe the report
+python -m nanochat.report reset
+
+# train tokenizer on ~1B characters
+python -m nanochat.dataset -n 4
+python -m scripts.tok_train --max_chars=1000000000
+python -m scripts.tok_eval
+
+# train a very small 4 layer model on the CPU
+# each optimization step processes a single sequence of 1024 tokens
+# we only run 50 steps of optimization (bump this to get better results)
+python -m scripts.base_train \
+    --depth=4 \
+    --max_seq_len=1024 \
+    --device_batch_size=1 \
+    --total_batch_size=1024 \
+    --eval_every=50 \
+    --eval_tokens=4096 \
+    --core_metric_every=50 \
+    --core_metric_max_per_task=12 \
+    --sample_every=50 \
+    --num_iterations=50
+python -m scripts.base_loss --device_batch_size=1 --split_tokens=4096
+python -m scripts.base_eval --max-per-task=16
+
+# midtraining
+python -m scripts.mid_train \
+    --max_seq_len=1024 \
+    --device_batch_size=1 \
+    --eval_every=50 \
+    --eval_tokens=4096 \
+    --total_batch_size=1024 \
+    --num_iterations=100
+# eval results will be terrible, this is just to execute the code paths.
+# note that we lower the execution memory limit to 1MB to avoid warnings on smaller systems
+python -m scripts.chat_eval --source=mid --max-new-tokens=128 --max-problems=20
+
+# SFT
+python -m scripts.chat_sft \
+    --device_batch_size=1 \
+    --target_examples_per_step=4 \
+    --num_iterations=100 \
+    --eval_steps=4 \
+    --eval_metrics_max_problems=16
+
+# Chat CLI
+# python -m scripts.chat_cli -p "Why is the sky blue?"
+
+# Chat Web
+# python -m scripts.chat_web
+
+python -m nanochat.report generate

nanochat/adamw.py

@@ -0,0 +1,76 @@
+"""
+Borrowed from modded-nanogpt. By Keller, @vagrawal, et al.
+Not a general optimizer! But works for our specific use.
+"""
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
+
+class DistAdamW(torch.optim.Optimizer):
+    """
+    Distributed AdamW optimizer.
+    In the style of ZeRO-2, i.e. sharded optimizer states and gradient reduction
+    """
+    def __init__(self, param_groups, lr: float = 1e-3, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.01):
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super().__init__(param_groups, defaults)
+
+    @torch.compile
+    @torch.no_grad()
+    def step(self):
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+        reduce_scatter_futures: list[torch.Future] = []
+        all_reduce_futures: list[torch.Future] = []
+        grad_slices = []
+        for group in self.param_groups:
+            params: list[Tensor] = group["params"]
+            for base_i in range(len(params)):
+                grad = params[base_i].grad
+                rank_size = grad.shape[0] // world_size
+                grad_slice = torch.empty_like(grad[:rank_size])
+                reduce_scatter_futures.append(dist.reduce_scatter_tensor(grad_slice, grad, op=dist.ReduceOp.AVG, async_op=True).get_future())
+                grad_slices.append(grad_slice)
+
+        idx = 0
+        for group in self.param_groups:
+            beta1, beta2 = group['betas']
+            eps = group['eps']
+            wd = group['weight_decay']
+            params = group['params']
+            for base in range(len(params)):
+                reduce_scatter_futures[idx].wait()
+                p = params[base]
+                rank_size = p.shape[0] // world_size
+                p_slice = p[rank * rank_size:(rank + 1) * rank_size]
+                lr = group['lr'] * getattr(p, "lr_mul", 1.0)
+                state = self.state[p]
+                g_slice = grad_slices[idx]
+                # State init
+                if not state:
+                    state['step'] = torch.tensor(0, dtype=torch.int64, device=p.device)
+                    state['exp_avg'] = torch.zeros_like(p_slice)
+                    state['exp_avg_sq'] = torch.zeros_like(p_slice)
+                exp_avg = state['exp_avg']
+                exp_avg_sq = state['exp_avg_sq']
+                state['step'] += 1
+                t = state['step']
+                # weight decay
+                if wd != 0:
+                    eff_weight_decay = lr * wd * getattr(p, "wd_mul", 1.0)
+                    p_slice.mul_(1 - eff_weight_decay)
+                # update running averages
+                exp_avg.mul_(beta1).add_(g_slice, alpha=1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(g_slice, g_slice, value=1 - beta2)
+                # bias corrections
+                bias1 = 1 - beta1 ** t
+                bias2 = 1 - beta2 ** t
+                # compute step
+                denom = exp_avg_sq.sqrt().add_(eps)
+                step_size = lr * (torch.sqrt(bias2) / bias1)
+                update = exp_avg.div(denom).mul_(step_size)
+                p_slice.add_(other=update, alpha=-1.0)
+                idx += 1
+                all_reduce_futures.append(dist.all_gather_into_tensor(p, p_slice, async_op=True).get_future())
+        torch.futures.collect_all(all_reduce_futures).wait()

nanochat/checkpoint_manager.py

@@ -0,0 +1,152 @@
+"""
+Utilities for saving and loading model/optim/state checkpoints.
+"""
+import os
+import re
+import glob
+import json
+import logging
+import torch
+
+from nanochat.common import get_base_dir
+from nanochat.gpt import GPT, GPTConfig
+from nanochat.tokenizer import get_tokenizer
+from nanochat.common import setup_default_logging
+
+# Set up logging
+setup_default_logging()
+logger = logging.getLogger(__name__)
+def log0(message):
+    if int(os.environ.get('RANK', 0)) == 0:
+        logger.info(message)
+
+def save_checkpoint(checkpoint_dir, step, model_data, optimizer_data, meta_data):
+    assert int(os.environ.get('RANK', 0)) == 0 # prevent footguns for now
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    # Save the model state (parameters)
+    model_path = os.path.join(checkpoint_dir, f"model_{step:06d}.pt")
+    torch.save(model_data, model_path)
+    log0(f"Saved model file to: {model_path}")
+    # Save the optimizer state (useful for SFT or any other fine-tuning)
+    if optimizer_data is not None:
+        optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}.pt")
+        torch.save(optimizer_data, optimizer_path)
+        log0(f"Saved optimizer file to: {optimizer_path}")
+    # Save the metadata dict as json
+    meta_path = os.path.join(checkpoint_dir, f"meta_{step:06d}.json")
+    with open(meta_path, "w", encoding="utf-8") as f:
+        json.dump(meta_data, f, indent=2)
+    log0(f"Saved metadata file to: {meta_path}")
+
+
+def load_checkpoint(checkpoint_dir, step, device, load_optimizer=False):
+    # Load the model state
+    model_path = os.path.join(checkpoint_dir, f"model_{step:06d}.pt")
+    model_data = torch.load(model_path, map_location=device)
+    # Load the optimizer state if requested
+    optimizer_data = None
+    if load_optimizer:
+        optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}.pt")
+        optimizer_data = torch.load(optimizer_path, map_location=device)
+    # Load the metadata
+    meta_path = os.path.join(checkpoint_dir, f"meta_{step:06d}.json")
+    with open(meta_path, "r", encoding="utf-8") as f:
+        meta_data = json.load(f)
+    return model_data, optimizer_data, meta_data
+
+
+def build_model(checkpoint_dir, step, device, phase):
+    """
+    A bunch of repetitive code to build a model from a given checkpoint.
+    Returns:
+    - base model - uncompiled, not wrapped in DDP
+    - tokenizer
+    - meta data saved during base model training
+    """
+    assert phase in ["train", "eval"], f"Invalid phase: {phase}"
+    model_data, optimizer_data, meta_data = load_checkpoint(checkpoint_dir, step, device, load_optimizer=False)
+    if device.type in {"cpu", "mps"}:
+        # Convert bfloat16 tensors to float for CPU inference
+        model_data = {
+            k: v.float() if v.dtype == torch.bfloat16 else v
+            for k, v in model_data.items()
+        }
+    # Hack: fix torch compile issue, which prepends all keys with _orig_mod.
+    model_data = {k.removeprefix("_orig_mod."): v for k, v in model_data.items()}
+    model_config_kwargs = meta_data["model_config"]
+    log0(f"Building model with config: {model_config_kwargs}")
+    model_config = GPTConfig(**model_config_kwargs)
+    with torch.device("meta"):
+        model = GPT(model_config)
+    # Load the model state
+    model.to_empty(device=device)
+    model.init_weights() # note: this is dumb, but we need to init the rotary embeddings. TODO: fix model re-init
+    model.load_state_dict(model_data, strict=True, assign=True)
+    # Put the model in the right training phase / mode
+    if phase == "eval":
+        model.eval()
+    else:
+        model.train()
+    # Load the Tokenizer
+    tokenizer = get_tokenizer()
+    # Sanity check: compatibility between model and tokenizer
+    assert tokenizer.get_vocab_size() == model_config_kwargs["vocab_size"]
+    return model, tokenizer, meta_data
+
+
+def find_largest_model(checkpoint_dir):
+    # attempt to guess the model tag: take the biggest model available
+    model_tags = [f for f in os.listdir(checkpoint_dir) if os.path.isdir(os.path.join(checkpoint_dir, f))]
+    if not model_tags:
+        raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}")
+    # 1) normally all model tags are of the form d<number>, try that first:
+    candidates = []
+    for model_tag in model_tags:
+        match = re.match(r"d(\d+)", model_tag)
+        if match:
+            model_depth = int(match.group(1))
+            candidates.append((model_depth, model_tag))
+    if candidates:
+        candidates.sort(key=lambda x: x[0], reverse=True)
+        return candidates[0][1]
+    # 2) if that failed, take the most recently updated model:
+    model_tags.sort(key=lambda x: os.path.getmtime(os.path.join(checkpoint_dir, x)), reverse=True)
+    return model_tags[0]
+
+
+def find_last_step(checkpoint_dir):
+    # Look into checkpoint_dir and find model_<step>.pt with the highest step
+    checkpoint_files = glob.glob(os.path.join(checkpoint_dir, "model_*.pt"))
+    if not checkpoint_files:
+        raise FileNotFoundError(f"No checkpoints found in {checkpoint_dir}")
+    last_step = int(max(os.path.basename(f).split("_")[-1].split(".")[0] for f in checkpoint_files))
+    return last_step
+
+# -----------------------------------------------------------------------------
+# convenience functions that take into account nanochat's directory structure
+
+def load_model_from_dir(checkpoints_dir, device, phase, model_tag=None, step=None):
+    if model_tag is None:
+        # guess the model tag by defaulting to the largest model
+        model_tag = find_largest_model(checkpoints_dir)
+        log0(f"No model tag provided, guessing model tag: {model_tag}")
+    checkpoint_dir = os.path.join(checkpoints_dir, model_tag)
+    if step is None:
+        # guess the step by defaulting to the last step
+        step = find_last_step(checkpoint_dir)
+    assert step is not None, f"No checkpoints found in {checkpoint_dir}"
+    # build the model
+    log0(f"Loading model from {checkpoint_dir} with step {step}")
+    model, tokenizer, meta_data = build_model(checkpoint_dir, step, device, phase)
+    return model, tokenizer, meta_data
+
+def load_model(source, *args, **kwargs):
+    model_dir = {
+        "base": "base_checkpoints",
+        "mid": "mid_checkpoints",
+        "sft": "chatsft_checkpoints",
+        "rl": "chatrl_checkpoints",
+    }[source]
+    base_dir = get_base_dir()
+    checkpoints_dir = os.path.join(base_dir, model_dir)
+    return load_model_from_dir(checkpoints_dir, *args, **kwargs)

nanochat/common.py

@@ -0,0 +1,188 @@
+"""
+Common utilities for nanochat.
+"""
+
+import os
+import re
+import logging
+import urllib.request
+import torch
+import torch.distributed as dist
+from filelock import FileLock
+
+class ColoredFormatter(logging.Formatter):
+    """Custom formatter that adds colors to log messages."""
+    # ANSI color codes
+    COLORS = {
+        'DEBUG': '\033[36m',    # Cyan
+        'INFO': '\033[32m',     # Green
+        'WARNING': '\033[33m',  # Yellow
+        'ERROR': '\033[31m',    # Red
+        'CRITICAL': '\033[35m', # Magenta
+    }
+    RESET = '\033[0m'
+    BOLD = '\033[1m'
+    def format(self, record):
+        # Add color to the level name
+        levelname = record.levelname
+        if levelname in self.COLORS:
+            record.levelname = f"{self.COLORS[levelname]}{self.BOLD}{levelname}{self.RESET}"
+        # Format the message
+        message = super().format(record)
+        # Add color to specific parts of the message
+        if levelname == 'INFO':
+            # Highlight numbers and percentages
+            message = re.sub(r'(\d+\.?\d*\s*(?:GB|MB|%|docs))', rf'{self.BOLD}\1{self.RESET}', message)
+            message = re.sub(r'(Shard \d+)', rf'{self.COLORS["INFO"]}{self.BOLD}\1{self.RESET}', message)
+        return message
+
+def setup_default_logging():
+    handler = logging.StreamHandler()
+    handler.setFormatter(ColoredFormatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+    logging.basicConfig(
+        level=logging.INFO,
+        handlers=[handler]
+    )
+
+setup_default_logging()
+logger = logging.getLogger(__name__)
+
+def get_base_dir():
+    # co-locate nanochat intermediates with other cached data in ~/.cache (by default)
+    if os.environ.get("NANOCHAT_BASE_DIR"):
+        nanochat_dir = os.environ.get("NANOCHAT_BASE_DIR")
+    else:
+        home_dir = os.path.expanduser("~")
+        cache_dir = os.path.join(home_dir, ".cache")
+        nanochat_dir = os.path.join(cache_dir, "nanochat")
+    os.makedirs(nanochat_dir, exist_ok=True)
+    return nanochat_dir
+
+def download_file_with_lock(url, filename, postprocess_fn=None):
+    """
+    Downloads a file from a URL to a local path in the base directory.
+    Uses a lock file to prevent concurrent downloads among multiple ranks.
+    """
+    base_dir = get_base_dir()
+    file_path = os.path.join(base_dir, filename)
+    lock_path = file_path + ".lock"
+
+    if os.path.exists(file_path):
+        return file_path
+
+    with FileLock(lock_path):
+        # Only a single rank can acquire this lock
+        # All other ranks block until it is released
+
+        # Recheck after acquiring lock
+        if os.path.exists(file_path):
+            return file_path
+
+        # Download the content as bytes
+        print(f"Downloading {url}...")
+        with urllib.request.urlopen(url) as response:
+            content = response.read() # bytes
+
+        # Write to local file
+        with open(file_path, 'wb') as f:
+            f.write(content)
+        print(f"Downloaded to {file_path}")
+
+        # Run the postprocess function if provided
+        if postprocess_fn is not None:
+            postprocess_fn(file_path)
+
+    return file_path
+
+def print0(s="",**kwargs):
+    ddp_rank = int(os.environ.get('RANK', 0))
+    if ddp_rank == 0:
+        print(s, **kwargs)
+
+def print_banner():
+    # Cool DOS Rebel font ASCII banner made with https://manytools.org/hacker-tools/ascii-banner/
+    banner = """
+                                                       █████                █████
+                                                      ░░███                ░░███
+     ████████    ██████   ████████    ██████   ██████  ░███████    ██████  ███████
+    ░░███░░███  ░░░░░███ ░░███░░███  ███░░███ ███░░███ ░███░░███  ░░░░░███░░░███░
+     ░███ ░███   ███████  ░███ ░███ ░███ ░███░███ ░░░  ░███ ░███   ███████  ░███
+     ░███ ░███  ███░░███  ░███ ░███ ░███ ░███░███  ███ ░███ ░███  ███░░███  ░███ ███
+     ████ █████░░████████ ████ █████░░██████ ░░██████  ████ █████░░███████  ░░█████
+    ░░░░ ░░░░░  ░░░░░░░░ ░░░░ ░░░░░  ░░░░░░   ░░░░░░  ░░░░ ░░░░░  ░░░░░░░░   ░░░░░
+    """
+    print0(banner)
+
+def is_ddp():
+    # TODO is there a proper way
+    return int(os.environ.get('RANK', -1)) != -1
+
+def get_dist_info():
+    if is_ddp():
+        assert all(var in os.environ for var in ['RANK', 'LOCAL_RANK', 'WORLD_SIZE'])
+        ddp_rank = int(os.environ['RANK'])
+        ddp_local_rank = int(os.environ['LOCAL_RANK'])
+        ddp_world_size = int(os.environ['WORLD_SIZE'])
+        return True, ddp_rank, ddp_local_rank, ddp_world_size
+    else:
+        return False, 0, 0, 1
+
+def autodetect_device_type():
+    # prefer to use CUDA if available, otherwise use MPS, otherwise fallback on CPU
+    if torch.cuda.is_available():
+        device_type = "cuda"
+    elif torch.backends.mps.is_available():
+        device_type = "mps"
+    else:
+        device_type = "cpu"
+    print0(f"Autodetected device type: {device_type}")
+    return device_type
+
+def compute_init(device_type="cuda"): # cuda|cpu|mps
+    """Basic initialization that we keep doing over and over, so make common."""
+
+    assert device_type in ["cuda", "mps", "cpu"], "Invalid device type atm"
+    if device_type == "cuda":
+        assert torch.cuda.is_available(), "Your PyTorch installation is not configured for CUDA but device_type is 'cuda'"
+    if device_type == "mps":
+        assert torch.backends.mps.is_available(), "Your PyTorch installation is not configured for MPS but device_type is 'mps'"
+
+    # Reproducibility
+    torch.manual_seed(42)
+    if device_type == "cuda":
+        torch.cuda.manual_seed(42)
+    # skipping full reproducibility for now, possibly investigate slowdown later
+    # torch.use_deterministic_algorithms(True)
+
+    # Precision
+    if device_type == "cuda":
+        torch.set_float32_matmul_precision("high") # uses tf32 instead of fp32 for matmuls
+
+    # Distributed setup: Distributed Data Parallel (DDP), optional, and requires CUDA
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
+    if ddp and device_type == "cuda":
+        device = torch.device("cuda", ddp_local_rank)
+        torch.cuda.set_device(device)  # make "cuda" default to this device
+        dist.init_process_group(backend="nccl", device_id=device)
+        dist.barrier()
+    else:
+        device = torch.device(device_type) # mps|cpu
+
+    if ddp_rank == 0:
+        logger.info(f"Distributed world size: {ddp_world_size}")
+
+    return ddp, ddp_rank, ddp_local_rank, ddp_world_size, device
+
+def compute_cleanup():
+    """Companion function to compute_init, to clean things up before script exit"""
+    if is_ddp():
+        dist.destroy_process_group()
+
+class DummyWandb:
+    """Useful if we wish to not use wandb but have all the same signatures"""
+    def __init__(self):
+        pass
+    def log(self, *args, **kwargs):
+        pass
+    def finish(self):
+        pass

nanochat/configurator.py

@@ -0,0 +1,56 @@
+"""
+Poor Man's Configurator. Probably a terrible idea. Example usage:
+$ python train.py config/override_file.py --batch_size=32
+this will first run config/override_file.py, then override batch_size to 32
+
+The code in this file will be run as follows from e.g. train.py:
+>>> exec(open('configurator.py').read())
+
+So it's not a Python module, it's just shuttling this code away from train.py
+The code in this script then overrides the globals()
+
+I know people are not going to love this, I just really dislike configuration
+complexity and having to prepend config. to every single variable. If someone
+comes up with a better simple Python solution I am all ears.
+"""
+
+import os
+import sys
+from ast import literal_eval
+
+def print0(s="",**kwargs):
+    ddp_rank = int(os.environ.get('RANK', 0))
+    if ddp_rank == 0:
+        print(s, **kwargs)
+
+for arg in sys.argv[1:]:
+    if '=' not in arg:
+        # assume it's the name of a config file
+        assert not arg.startswith('--')
+        config_file = arg
+        print0(f"Overriding config with {config_file}:")
+        with open(config_file) as f:
+            print0(f.read())
+        exec(open(config_file).read())
+    else:
+        # assume it's a --key=value argument
+        assert arg.startswith('--')
+        key, val = arg.split('=')
+        key = key[2:]
+        if key in globals():
+            try:
+                # attempt to eval it it (e.g. if bool, number, or etc)
+                attempt = literal_eval(val)
+            except (SyntaxError, ValueError):
+                # if that goes wrong, just use the string
+                attempt = val
+            # ensure the types match ok
+            if globals()[key] is not None:
+                attempt_type = type(attempt)
+                default_type = type(globals()[key])
+                assert attempt_type == default_type, f"Type mismatch: {attempt_type} != {default_type}"
+            # cross fingers
+            print0(f"Overriding: {key} = {attempt}")
+            globals()[key] = attempt
+        else:
+            raise ValueError(f"Unknown config key: {key}")

nanochat/core_eval.py

@@ -0,0 +1,262 @@
+"""
+Functions for evaluating the CORE metric, as described in the DCLM paper.
+https://arxiv.org/abs/2406.11794
+
+TODOs:
+- All tasks ~match except for squad. We get 31% reference is 37%. Figure out why.
+"""
+import random
+
+from jinja2 import Template
+import torch
+import torch.distributed as dist
+
+# -----------------------------------------------------------------------------
+# Prompt rendering utilities
+
+def render_prompts_mc(item, continuation_delimiter, fewshot_examples=None):
+    """Render complete prompts for a multiple choice question"""
+    template_str = """
+{%- for example in fewshot_examples -%}
+{{ example.query }}{{ continuation_delimiter }}{{ example.choices[example.gold] }}
+
+{% endfor -%}
+{{ item.query }}{{ continuation_delimiter }}{{ choice }}""".strip()
+    template = Template(template_str)
+    fewshot_examples = fewshot_examples or []
+    context = {
+        'fewshot_examples': fewshot_examples,
+        'continuation_delimiter': continuation_delimiter,
+        'item': item
+    }
+    prompts = [template.render(choice=choice, **context) for choice in item['choices']]
+    return prompts
+
+
+def render_prompts_schema(item, continuation_delimiter, fewshot_examples=None):
+    """Render complete prompts for a schema question"""
+    template_str = """
+{%- for example in fewshot_examples -%}
+{{ example.context_options[example.gold] }}{{ continuation_delimiter }}{{ example.continuation }}
+
+{% endfor -%}
+{{ context }}{{ continuation_delimiter }}{{ item.continuation }}""".strip()
+    template = Template(template_str)
+    fewshot_examples = fewshot_examples or []
+    context = {
+        'fewshot_examples': fewshot_examples,
+        'continuation_delimiter': continuation_delimiter,
+        'item': item
+    }
+    prompts = [template.render(context=context_option, **context)
+               for context_option in item['context_options']]
+    return prompts
+
+
+def render_prompts_lm(item, continuation_delimiter, fewshot_examples=None):
+    """
+    Render complete prompt for a language modeling task.
+    Notice that we manually trim the context in the template,
+    which in some datasets seems to have trailing whitespace (which we don't want).
+    """
+    template_str = """
+{%- for example in fewshot_examples -%}
+{{ example.context | trim }}{{ continuation_delimiter }}{{ example.continuation }}
+
+{% endfor -%}
+{{ item.context | trim }}{{ continuation_delimiter }}{% if include_continuation %}{{ item.continuation }}{% endif %}""".strip()
+    template = Template(template_str)
+    fewshot_examples = fewshot_examples or []
+    context = {
+        'fewshot_examples': fewshot_examples,
+        'continuation_delimiter': continuation_delimiter,
+        'item': item
+    }
+    # Return two prompts: without and with the continuation
+    prompt_without = template.render(include_continuation=False, **context)
+    prompt_with = template.render(include_continuation=True, **context)
+    # Due to the way the data seems to be stored, I think I need to strip in the case of LM here.
+    # Otherwise we may get trailing whitespaces in prompt_without (which get absorbed into the next
+    # token in prompt_with), meaning we don't get a nice and clean prefix in the token space
+    # to detect the final continuation. Tokenizers...
+    prompt_without = prompt_without.strip()
+    return [prompt_without, prompt_with]
+
+
+def find_common_length(token_sequences, direction='left'):
+    """
+    Find the length of the common prefix or suffix across token sequences
+    - direction: 'left' for prefix, 'right' for suffix
+    """
+    min_len = min(len(seq) for seq in token_sequences)
+    indices = {
+        'left': range(min_len),
+        'right': range(-1, -min_len-1, -1)
+    }[direction]
+    # Find the first position where the token sequences differ
+    for i, idx in enumerate(indices):
+        token = token_sequences[0][idx]
+        if not all(seq[idx] == token for seq in token_sequences):
+            return i
+    return min_len
+
+
+def stack_sequences(tokens, pad_token_id):
+    """Stack up a list of token sequences, pad to longest on the right"""
+    bsz, seq_len = len(tokens), max(len(x) for x in tokens)
+    input_ids = torch.full((bsz, seq_len), pad_token_id, dtype=torch.long)
+    for i, x in enumerate(tokens):
+        input_ids[i, :len(x)] = torch.tensor(x, dtype=torch.long)
+    return input_ids
+
+
+def batch_sequences_mc(tokenizer, prompts):
+    # In multiple choice, contexts are the same but the continuation is different (common prefix)
+    tokens = tokenizer(prompts, prepend=tokenizer.get_bos_token_id())
+    # figure out the start and end of each continuation
+    answer_start_idx = find_common_length(tokens, direction='left')
+    start_indices = [answer_start_idx] * len(prompts)
+    end_indices = [len(x) for x in tokens]
+    return tokens, start_indices, end_indices
+
+
+def batch_sequences_schema(tokenizer, prompts):
+    # In schema tasks, contexts vary but continuation is the same (common suffix)
+    tokens = tokenizer(prompts, prepend=tokenizer.get_bos_token_id())
+    # figure out the start and end of each context
+    suffix_length = find_common_length(tokens, direction='right')
+    end_indices = [len(x) for x in tokens]
+    start_indices = [ei - suffix_length for ei in end_indices]
+    return tokens, start_indices, end_indices
+
+
+def batch_sequences_lm(tokenizer, prompts):
+    # In LM tasks, we have two prompts: without and with continuation
+    tokens = tokenizer(prompts, prepend=tokenizer.get_bos_token_id())
+    tokens_without, tokens_with = tokens
+    start_idx, end_idx = len(tokens_without), len(tokens_with)
+    assert start_idx < end_idx, "prompt without is supposed to be a prefix of prompt with"
+    assert tokens_without == tokens_with[:start_idx], "prompt without is supposed to be a prefix of prompt with"
+    # we only need the with continuation prompt in the LM task, i.e. batch size of 1
+    return [tokens_with], [start_idx], [end_idx]
+
+
+@torch.no_grad()
+def forward_model(model, input_ids):
+    """
+    Take BxT tensor of token ids, return BxT tensor of losses and argmax predictions.
+    The last column of losses is set to nan because we don't have autoregressive targets there.
+    """
+    batch_size, seq_len = input_ids.size()
+    outputs = model(input_ids)
+    # Roll the tensor to the left by one position to get the (autoregressive) target ids
+    target_ids = torch.roll(input_ids, shifts=-1, dims=1)
+    # Calculate cross entropy at all positions
+    losses = torch.nn.functional.cross_entropy(
+        outputs.view(batch_size * seq_len, -1),
+        target_ids.view(batch_size * seq_len),
+        reduction='none'
+    ).view(batch_size, seq_len)
+    # Set the last column to be nan because there is no autoregressive loss there
+    losses[:, -1] = float('nan')
+    # Get the argmax predictions at each position
+    predictions = outputs.argmax(dim=-1)
+    return losses, predictions
+
+
+@torch.no_grad()
+def evaluate_example(idx, model, tokenizer, data, device, task_meta):
+    """Evaluate a single example, return True if correct, False otherwise"""
+    item = data[idx]
+    task_type = task_meta['task_type']
+    num_fewshot = task_meta['num_fewshot']
+    continuation_delimiter = task_meta['continuation_delimiter']
+
+    # Sample few-shot examples (excluding current item)
+    fewshot_examples = []
+    if num_fewshot > 0:
+        rng = random.Random(1234 + idx)
+        available_indices = [i for i in range(len(data)) if i != idx]
+        fewshot_indices = rng.sample(available_indices, num_fewshot)
+        fewshot_examples = [data[i] for i in fewshot_indices]
+
+    # Render prompts and batch sequences based on task type
+    if task_type == 'multiple_choice':
+        prompts = render_prompts_mc(item, continuation_delimiter, fewshot_examples)
+        tokens, start_idxs, end_idxs = batch_sequences_mc(tokenizer, prompts)
+    elif task_type == 'schema':
+        prompts = render_prompts_schema(item, continuation_delimiter, fewshot_examples)
+        tokens, start_idxs, end_idxs = batch_sequences_schema(tokenizer, prompts)
+    elif task_type == 'language_modeling':
+        prompts = render_prompts_lm(item, continuation_delimiter, fewshot_examples)
+        tokens, start_idxs, end_idxs = batch_sequences_lm(tokenizer, prompts)
+    else:
+        raise ValueError(f"Unsupported task type: {task_type}")
+
+    # Some models can't forward sequences beyond a certain length (e.g. GPT-2)
+    # In these cases, we have to truncate sequences to max length and adjust the indices
+    if hasattr(model, 'max_seq_len') and model.max_seq_len is not None:
+        max_tokens = model.max_seq_len
+        new_tokens, new_start_idxs, new_end_idxs = [], [], []
+        for t, s, e in zip(tokens, start_idxs, end_idxs):
+            if len(t) > max_tokens:
+                num_to_crop = len(t) - max_tokens
+                new_tokens.append(t[-max_tokens:]) # take the last max_tokens tokens
+                new_start_idxs.append(s - num_to_crop) # shift the indices down
+                new_end_idxs.append(e - num_to_crop)
+                assert s - num_to_crop >= 0, "this should never happen right?"
+                assert e - num_to_crop >= 0, "this should never happen right?"
+            else:
+                new_tokens.append(t) # keep unchanged
+                new_start_idxs.append(s)
+                new_end_idxs.append(e)
+        tokens, start_idxs, end_idxs = new_tokens, new_start_idxs, new_end_idxs
+
+    # Stack up all the sequences into a batch
+    pad_token_id = tokenizer.get_bos_token_id() # use BOS as pad token is ok
+    input_ids = stack_sequences(tokens, pad_token_id)
+    input_ids = input_ids.to(device)
+
+    # Forward the model, get the autoregressive loss and argmax prediction at each token
+    losses, predictions = forward_model(model, input_ids)
+
+    # See if the losses/predictions come out correctly
+    if task_type == 'language_modeling':
+        # language modeling task is currently always batch size 1
+        si = start_idxs[0]
+        ei = end_idxs[0]
+        # predictions[i] predict input_ids[i+1] autoregressively
+        predicted_tokens = predictions[0, si-1:ei-1]
+        actual_tokens = input_ids[0, si:ei]
+        is_correct = torch.all(predicted_tokens == actual_tokens).item()
+    elif task_type in ['multiple_choice', 'schema']:
+        # For MC/schema: find the option with lowest average loss
+        mean_losses = [losses[i, si-1:ei-1].mean().item()
+                        for i, (si, ei) in enumerate(zip(start_idxs, end_idxs))]
+        pred_idx = mean_losses.index(min(mean_losses))
+        is_correct = pred_idx == item['gold']
+    else:
+        raise ValueError(f"Unsupported task type: {task_type}")
+
+    return is_correct
+
+
+def evaluate_task(model, tokenizer, data, device, task_meta):
+    """
+    This function is responsible for evaluating one task across many examples.
+    It also handles dispatch to all processes if the script is run with torchrun.
+    """
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    correct = torch.zeros(len(data), dtype=torch.float32, device=device)
+    # stride the examples to each rank
+    for idx in range(rank, len(data), world_size):
+        is_correct = evaluate_example(idx, model, tokenizer, data, device, task_meta)
+        correct[idx] = float(is_correct)
+    # sync results across all the processes if running distributed
+    if world_size > 1:
+        dist.barrier()
+        dist.all_reduce(correct, op=dist.ReduceOp.SUM)
+    # compute the mean
+    mean_correct = correct.mean().item()
+    return mean_correct

nanochat/dataloader.py

@@ -0,0 +1,49 @@
+from collections import deque
+
+import torch
+
+from nanochat.common import get_dist_info
+from nanochat.dataset import parquets_iter_batched
+from nanochat.tokenizer import get_tokenizer
+
+def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokenizer_batch_size=128, device="cuda"):
+    """Stream pretraining text from parquet files, tokenize, yield training batches."""
+    assert split in ["train", "val"], "split must be 'train' or 'val'"
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
+    needed_tokens = B * T + 1 # +1 is because we also need the target at the last token
+    # get the tokenizer and the bos token
+    tokenizer = get_tokenizer()
+    bos_token = tokenizer.get_bos_token_id()
+    # scratch buffer holds the tokens for one iteration
+    token_buffer = deque() # we stream tokens on the right and pop from the left
+
+    # infinite iterator over document batches
+    def document_batches():
+        while True:
+            # batch will iterate in group size of the parquet files, usually e.g. 1024 rows
+            for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size):
+                # for the tokenizer we might want to go in usually smaller batches, e.g. 128 rows
+                for i in range(0, len(batch), tokenizer_batch_size):
+                    yield batch[i:i+tokenizer_batch_size]
+    batches = document_batches()
+
+    batch_index = 0
+    while True:
+        # Accumulate enough tokens for one iteration before yielding.
+        while len(token_buffer) < needed_tokens:
+            doc_batch = next(batches)
+            token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads)
+            for tokens in token_lists:
+                token_buffer.extend(tokens)
+            batch_index += 1
+        # Move tokens from the deque into the scratch buffer
+        tokens = [token_buffer.popleft() for _ in range(needed_tokens)]
+        # CUDA supports memory pinning for faster transfers between CPU and GPU:
+        scratch = torch.tensor(tokens, dtype=torch.int64, pin_memory=(device == "cuda"))
+        # Create the inputs/targets as 1D tensors
+        inputs_cpu = scratch[:-1].to(dtype=torch.int32)
+        targets_cpu = scratch[1:]
+        # Reshape to 2D and move to GPU async
+        inputs = inputs_cpu.view(B, T).to(device=device, dtype=torch.int32, non_blocking=True)
+        targets = targets_cpu.view(B, T).to(device=device, dtype=torch.int64, non_blocking=True)
+        yield inputs, targets

nanochat/dataset.py

@@ -0,0 +1,128 @@
+"""
+The base/pretraining dataset is a set of parquet files.
+This file contains utilities for:
+- iterating over the parquet files and yielding documents from it
+- download the files on demand if they are not on disk
+
+For details of how the dataset was prepared, see `repackage_data_reference.py`.
+"""
+
+import os
+import argparse
+import time
+import requests
+import pyarrow.parquet as pq
+from multiprocessing import Pool
+
+from nanochat.common import get_base_dir
+
+# -----------------------------------------------------------------------------
+# The specifics of the current pretraining dataset
+
+# The URL on the internet where the data is hosted and downloaded from on demand
+BASE_URL = "https://huggingface.co/datasets/karpathy/fineweb-edu-100b-shuffle/resolve/main"
+MAX_SHARD = 1822 # the last datashard is shard_01822.parquet
+index_to_filename = lambda index: f"shard_{index:05d}.parquet" # format of the filenames
+base_dir = get_base_dir()
+DATA_DIR = os.path.join(base_dir, "base_data")
+os.makedirs(DATA_DIR, exist_ok=True)
+
+# -----------------------------------------------------------------------------
+# These functions are useful utilities to other modules, can/should be imported
+
+def list_parquet_files(data_dir=None):
+    """ Looks into a data dir and returns full paths to all parquet files. """
+    data_dir = DATA_DIR if data_dir is None else data_dir
+    parquet_files = sorted([
+        f for f in os.listdir(data_dir)
+        if f.endswith('.parquet') and not f.endswith('.tmp')
+    ])
+    parquet_paths = [os.path.join(data_dir, f) for f in parquet_files]
+    return parquet_paths
+
+def parquets_iter_batched(split, start=0, step=1):
+    """
+    Iterate through the dataset, in batches of underlying row_groups for efficiency.
+    - split can be "train" or "val". the last parquet file will be val.
+    - start/step are useful for skipping rows in DDP. e.g. start=rank, step=world_size
+    """
+    assert split in ["train", "val"], "split must be 'train' or 'val'"
+    parquet_paths = list_parquet_files()
+    parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:]
+    for filepath in parquet_paths:
+        pf = pq.ParquetFile(filepath)
+        for rg_idx in range(start, pf.num_row_groups, step):
+            rg = pf.read_row_group(rg_idx)
+            texts = rg.column('text').to_pylist()
+            yield texts
+
+# -----------------------------------------------------------------------------
+def download_single_file(index):
+    """ Downloads a single file index, with some backoff """
+
+    # Construct the local filepath for this file and skip if it already exists
+    filename = index_to_filename(index)
+    filepath = os.path.join(DATA_DIR, filename)
+    if os.path.exists(filepath):
+        print(f"Skipping {filepath} (already exists)")
+        return True
+
+    # Construct the remote URL for this file
+    url = f"{BASE_URL}/{filename}"
+    print(f"Downloading {filename}...")
+
+    # Download with retries
+    max_attempts = 5
+    for attempt in range(1, max_attempts + 1):
+        try:
+            response = requests.get(url, stream=True, timeout=30)
+            response.raise_for_status()
+            # Write to temporary file first
+            temp_path = filepath + f".tmp"
+            with open(temp_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=1024 * 1024):  # 1MB chunks
+                    if chunk:
+                        f.write(chunk)
+            # Move temp file to final location
+            os.rename(temp_path, filepath)
+            print(f"Successfully downloaded {filename}")
+            return True
+
+        except (requests.RequestException, IOError) as e:
+            print(f"Attempt {attempt}/{max_attempts} failed for {filename}: {e}")
+            # Clean up any partial files
+            for path in [filepath + f".tmp", filepath]:
+                if os.path.exists(path):
+                    try:
+                        os.remove(path)
+                    except:
+                        pass
+            # Try a few times with exponential backoff: 2^attempt seconds
+            if attempt < max_attempts:
+                wait_time = 2 ** attempt
+                print(f"Waiting {wait_time} seconds before retry...")
+                time.sleep(wait_time)
+            else:
+                print(f"Failed to download {filename} after {max_attempts} attempts")
+                return False
+
+    return False
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download FineWeb-Edu 100BT dataset shards")
+    parser.add_argument("-n", "--num-files", type=int, default=-1, help="Number of shards to download (default: -1), -1 = disable")
+    parser.add_argument("-w", "--num-workers", type=int, default=4, help="Number of parallel download workers (default: 4)")
+    args = parser.parse_args()
+
+    num = MAX_SHARD + 1 if args.num_files == -1 else min(args.num_files, MAX_SHARD + 1)
+    ids_to_download = list(range(num))
+    print(f"Downloading {len(ids_to_download)} shards using {args.num_workers} workers...")
+    print(f"Target directory: {DATA_DIR}")
+    print()
+    with Pool(processes=args.num_workers) as pool:
+        results = pool.map(download_single_file, ids_to_download)
+
+    # Report results
+    successful = sum(1 for success in results if success)
+    print(f"Done! Downloaded: {successful}/{len(ids_to_download)} shards to {DATA_DIR}")

nanochat/engine.py

@@ -0,0 +1,371 @@
+"""
+Engine for efficient inference of our models.
+
+Everything works around token sequences:
+- The user can send token sequences to the engine
+- The engine returns the next token
+
+Notes:
+- The engine knows nothing about tokenization, it's purely token id sequences.
+
+The whole thing is made as efficient as possible.
+"""
+
+import torch
+import torch.nn.functional as F
+import signal
+import warnings
+from contextlib import contextmanager
+from collections import deque
+from nanochat.common import compute_init
+from nanochat.checkpoint_manager import load_model
+
+# -----------------------------------------------------------------------------
+# Calculator tool helpers
+@contextmanager
+def timeout(duration, formula):
+    def timeout_handler(signum, frame):
+        raise Exception(f"'{formula}': timed out after {duration} seconds")
+
+    signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(duration)
+    yield
+    signal.alarm(0)
+
+def eval_with_timeout(formula, max_time=3):
+    try:
+        with timeout(max_time, formula):
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", SyntaxWarning)
+                return eval(formula, {"__builtins__": {}}, {})
+    except Exception as e:
+        signal.alarm(0)
+        # print(f"Warning: Failed to eval {formula}, exception: {e}") # it's ok ignore wrong calculator usage
+        return None
+
+def use_calculator(expr):
+    """
+    Evaluate a Python expression safely.
+    Supports both math expressions and string operations like .count()
+    """
+    # Remove commas from numbers
+    expr = expr.replace(",", "")
+
+    # Check if it's a pure math expression (old behavior)
+    if all([x in "0123456789*+-/.() " for x in expr]):
+        if "**" in expr:  # disallow power operator
+            return None
+        return eval_with_timeout(expr)
+
+    # Check if it's a string operation we support
+    # Allow: strings (single/double quotes), .count(), letters, numbers, spaces, parens
+    allowed_chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'\"()._ "
+    if not all([x in allowed_chars for x in expr]):
+        return None
+
+    # Disallow dangerous patterns
+    dangerous_patterns = ['__', 'import', 'exec', 'eval', 'compile', 'open', 'file',
+                         'input', 'raw_input', 'globals', 'locals', 'vars', 'dir',
+                         'getattr', 'setattr', 'delattr', 'hasattr']
+    expr_lower = expr.lower()
+    if any(pattern in expr_lower for pattern in dangerous_patterns):
+        return None
+
+    # Only allow .count() method for now (can expand later)
+    if '.count(' not in expr:
+        return None
+
+    # Evaluate with timeout
+    return eval_with_timeout(expr)
+
+# -----------------------------------------------------------------------------
+class KVCache:
+    """
+    Works hand-in-hand with the GPT model to maintain the KV cache.
+    Note that the .pos advances automatically after the last layer of the Transformer inserts.
+    """
+
+    def __init__(self, batch_size, num_heads, seq_len, head_dim, num_layers):
+        # Each of K/V is of shape (B, H, T, D) and we have one per layer of the Transformer.
+        self.kv_shape = (num_layers, 2, batch_size, num_heads, seq_len, head_dim)
+        self.kv_cache = None
+        self.pos = 0 # current position in time in the cache
+
+    def reset(self):
+        self.pos = 0
+
+    def get_pos(self):
+        return self.pos
+
+    def prefill(self, other):
+        """
+        Prefill given another KV cache. Optionally expand along batch dim.
+        This is used when we do batch 1 prefill and then want to generate
+        multiple samples in parallel from there.
+        """
+        # 1) validate the shapes
+        assert self.kv_cache is None, "Cannot prefill a non-empty KV cache"
+        assert other.kv_cache is not None, "Cannot prefill with a None KV cache"
+        for ix, (dim1, dim2) in enumerate(zip(self.kv_shape, other.kv_shape)):
+            if ix in [0, 1, 3, 5]:
+                # num_layers, batch_size, num_heads, head_dim must match
+                assert dim1 == dim2, f"Dim {ix} mismatch: {dim1} != {dim2}"
+            elif ix == 2:
+                # batch_size can be expanded
+                assert dim1 == dim2 or dim2 == 1, f"Batch dim mismatch: {dim1} != {dim2}"
+            elif ix == 4:
+                # seq_len: self must be longer than other
+                assert dim1 >= dim2, f"Seq len mismatch: {dim1} < {dim2}"
+        # 2) initialize the cache
+        dtype, device = other.kv_cache.dtype, other.kv_cache.device
+        self.kv_cache = torch.empty(self.kv_shape, dtype=dtype, device=device)
+        # 3) copy the data over
+        self.kv_cache[:, :, :, :, :other.pos, :] = other.kv_cache
+        # 4) update the pos
+        self.pos = other.pos
+
+    def insert_kv(self, layer_idx, k, v):
+        # Lazy initialize the cache here because we need to know the dtype/device
+        if self.kv_cache is None:
+            self.kv_cache = torch.empty(self.kv_shape, dtype=k.dtype, device=k.device)
+        # Insert new keys/values to the cache and return the full cache so far
+        B, H, T_add, D = k.size()
+        t0, t1 = self.pos, self.pos + T_add
+        # Dynamically grow the cache if needed
+        if t1 > self.kv_cache.size(4):
+            t_needed = t1 + 1024 # as much as we need plus buffer of 1024
+            t_needed = (t_needed + 1023) & ~1023 # then round up to the nearest multiple of 1024
+            additional_shape = list(self.kv_cache.shape)
+            additional_shape[4] = t_needed - self.kv_cache.size(4)
+            additional_cache = torch.empty(additional_shape, dtype=k.dtype, device=k.device)
+            self.kv_cache = torch.cat([self.kv_cache, additional_cache], dim=4).contiguous()
+            self.kv_shape = self.kv_cache.shape
+        # Insert k, v into the cache
+        self.kv_cache[layer_idx, 0, :, :, t0:t1] = k
+        self.kv_cache[layer_idx, 1, :, :, t0:t1] = v
+        # Return the full cached keys/values up to current position (as a view)
+        key_view = self.kv_cache[layer_idx, 0, :, :, :t1]
+        value_view = self.kv_cache[layer_idx, 1, :, :, :t1]
+        # Increment pos after the last layer of the Transformer processes
+        if layer_idx == self.kv_cache.size(0) - 1:
+            self.pos = t1
+        return key_view, value_view
+
+
+# -----------------------------------------------------------------------------
+@torch.inference_mode()
+def sample_next_token(logits, rng, temperature=1.0, top_k=None):
+    """Sample a single next token from given logits of shape (B, vocab_size). Returns (B, 1)."""
+    assert temperature >= 0.0, "temperature must be non-negative"
+    if temperature == 0.0:
+        return torch.argmax(logits, dim=-1, keepdim=True)
+    if top_k is not None:
+        k = min(top_k, logits.size(-1))
+        vals, idx = torch.topk(logits, k, dim=-1)
+        vals = vals / temperature
+        probs = F.softmax(vals, dim=-1)
+        choice = torch.multinomial(probs, num_samples=1, generator=rng)
+        return idx.gather(1, choice)
+    else:
+        logits = logits / temperature
+        probs = F.softmax(logits, dim=-1)
+        return torch.multinomial(probs, num_samples=1, generator=rng)
+
+# -----------------------------------------------------------------------------
+
+class RowState:
+    # Per-row state tracking during generation
+    def __init__(self, current_tokens=None):
+        self.current_tokens = current_tokens or [] # Current token sequence for this row
+        self.forced_tokens = deque() # Queue of tokens to force inject
+        self.in_python_block = False # Whether we are inside a python block
+        self.python_expr_tokens = [] # Tokens of the current python expression
+        self.completed = False # Whether this row has completed generation
+
+class Engine:
+
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer # needed for tool use
+
+    @torch.inference_mode()
+    def generate(self, tokens, num_samples=1, max_tokens=None, temperature=1.0, top_k=None, seed=42):
+        """Same as generate, but does single prefill and then clones the KV cache."""
+        assert isinstance(tokens, list) and isinstance(tokens[0], int), "expecting list of ints"
+        device = self.model.get_device()
+        rng = torch.Generator(device=device)
+        rng.manual_seed(seed)
+
+        # Get the special tokens we need to coordinate the tool use state machine
+        get_special = lambda s: self.tokenizer.encode_special(s)
+        python_start = get_special("<|python_start|>")
+        python_end = get_special("<|python_end|>")
+        output_start = get_special("<|output_start|>")
+        output_end = get_special("<|output_end|>")
+        assistant_end = get_special("<|assistant_end|>") # if sampled, ends row
+        bos = self.tokenizer.get_bos_token_id() # if sampled, ends row
+
+        # 1) Run a batch 1 prefill of the prompt tokens
+        m = self.model.config
+        kv_model_kwargs = {"num_heads": m.n_kv_head, "head_dim": m.n_embd // m.n_head, "num_layers": m.n_layer}
+        kv_cache_prefill = KVCache(
+            batch_size=1,
+            seq_len=len(tokens),
+            **kv_model_kwargs,
+        )
+        ids = torch.tensor([tokens], dtype=torch.long, device=device)
+        logits = self.model.forward(ids, kv_cache=kv_cache_prefill)
+        logits = logits[:, -1, :]
+        next_ids = sample_next_token(logits, rng, temperature, top_k)  # (B, 1)
+        sampled_tokens = next_ids[:, 0].tolist()
+
+        # 2) Replicate the KV cache for each sample/row
+        kv_length_hint = (len(tokens) + max_tokens) if max_tokens is not None else self.model.config.sequence_len
+        kv_cache_decode = KVCache(
+            batch_size=num_samples,
+            seq_len=kv_length_hint,
+            **kv_model_kwargs,
+        )
+        kv_cache_decode.prefill(kv_cache_prefill)
+        del kv_cache_prefill # no need to keep this memory around
+
+        # 3) Initialize states for each sample
+        row_states = [RowState(tokens.copy()) for _ in range(num_samples)]
+
+        # 4) Main generation loop
+        num_generated = 0
+        first_iteration = True
+        while True:
+            # Stop condition: we've reached max tokens
+            if max_tokens is not None and num_generated >= max_tokens:
+                break
+            # Stop condition: all rows are completed
+            if all(state.completed for state in row_states):
+                break
+
+            # Get sampled tokens - either from prefill or from forward pass
+            if first_iteration:
+                # Use the tokens we already sampled from prefill
+                sampled_tokens = [sampled_tokens[0]] * num_samples  # Broadcast first token to all rows
+                # TODO: we should sample a token for each row instead of broadcasting
+                first_iteration = False
+            else:
+                # Forward the model and get the next token for each row
+                logits = self.model.forward(ids, kv_cache=kv_cache_decode)  # (B, T, vocab_size)
+                logits = logits[:, -1, :]  # (B, vocab_size) at last time step
+                next_ids = sample_next_token(logits, rng, temperature, top_k)  # (B, 1)
+                sampled_tokens = next_ids[:, 0].tolist()
+
+            # Process each row: choose the next token, update state, optional tool use
+            token_column = [] # contains the next token id along each row
+            token_masks = [] # contains the mask (was it sampled (1) or forced (0)?) along each row
+            for i, state in enumerate(row_states):
+                # Select the next token in this row
+                is_forced = len(state.forced_tokens) > 0 # are there tokens waiting to be forced in deque?
+                token_masks.append(0 if is_forced else 1) # mask is 0 if forced, 1 if sampled
+                next_token = state.forced_tokens.popleft() if is_forced else sampled_tokens[i]
+                token_column.append(next_token)
+                # Update the state of this row to include the next token
+                state.current_tokens.append(next_token)
+                # On <|assistant_end|> or <|bos|>, mark the row as completed
+                if next_token == assistant_end or next_token == bos:
+                    state.completed = True
+                # Handle tool logic
+                if next_token == python_start:
+                    state.in_python_block = True
+                    state.python_expr_tokens = []
+                elif next_token == python_end and state.in_python_block:
+                    state.in_python_block = False
+                    if state.python_expr_tokens:
+                        expr = self.tokenizer.decode(state.python_expr_tokens)
+                        result = use_calculator(expr)
+                        if result is not None:
+                            result_tokens = self.tokenizer.encode(str(result))
+                            state.forced_tokens.append(output_start)
+                            state.forced_tokens.extend(result_tokens)
+                            state.forced_tokens.append(output_end)
+                    state.python_expr_tokens = []
+                elif state.in_python_block:
+                    state.python_expr_tokens.append(next_token)
+
+            # Yield the token column
+            yield token_column, token_masks
+            num_generated += 1
+            # Prepare ids for next iteration
+            ids = torch.tensor(token_column, dtype=torch.long, device=device).unsqueeze(1)
+
+    def generate_batch(self, tokens, num_samples=1, **kwargs):
+        """
+        Non-streaming batch generation that just returns the final token sequences.
+        Returns a list of token sequences (list of lists of ints).
+        Terminal tokens (assistant_end, bos) are not included in the results.
+        """
+        assistant_end = self.tokenizer.encode_special("<|assistant_end|>")
+        bos = self.tokenizer.get_bos_token_id()
+        results = [tokens.copy() for _ in range(num_samples)]
+        masks = [[0] * len(tokens) for _ in range(num_samples)]
+        completed = [False] * num_samples
+        for token_column, token_masks in self.generate(tokens, num_samples, **kwargs):
+            for i, (token, mask) in enumerate(zip(token_column, token_masks)):
+                if not completed[i]:
+                    if token == assistant_end or token == bos:
+                        completed[i] = True
+                    else:
+                        results[i].append(token)
+                        masks[i].append(mask)
+            # Stop if all rows are completed
+            if all(completed):
+                break
+        return results, masks
+
+
+if __name__ == "__main__":
+    """
+    Quick inline test to make sure that the naive/slow model.generate function
+    is equivalent to the faster Engine.generate function here.
+    """
+    import time
+    # init compute
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init()
+    # load the model and tokenizer
+    model, tokenizer, meta = load_model("base", device, phase="eval")
+    bos_token_id = tokenizer.get_bos_token_id()
+    # common hyperparameters
+    kwargs = dict(max_tokens=64, temperature=0.0)
+    # set the starting prompt
+    prompt_tokens = tokenizer.encode("The chemical formula of water is", prepend=bos_token_id)
+    # generate the reference sequence using the model.generate() function
+    generated_tokens = []
+    torch.cuda.synchronize()
+    t0 = time.time()
+    stream = model.generate(prompt_tokens, **kwargs)
+    for token in stream:
+        generated_tokens.append(token)
+        chunk = tokenizer.decode([token])
+        print(chunk, end="", flush=True)
+    print()
+    torch.cuda.synchronize()
+    t1 = time.time()
+    print(f"Reference time: {t1 - t0:.2f}s")
+    reference_ids = generated_tokens
+    # generate tokens with Engine
+    generated_tokens = []
+    engine = Engine(model, tokenizer)
+    stream = engine.generate(prompt_tokens, num_samples=1, **kwargs) # note: runs in fp32
+    torch.cuda.synchronize()
+    t0 = time.time()
+    for token_column, token_masks in stream:
+        token = token_column[0] # only print out the first row
+        generated_tokens.append(token)
+        chunk = tokenizer.decode([token])
+        print(chunk, end="", flush=True)
+    print()
+    torch.cuda.synchronize()
+    t1 = time.time()
+    print(f"Engine time: {t1 - t0:.2f}s")
+    # compare the two sequences
+    for i in range(len(reference_ids)):
+        if reference_ids[i] != generated_tokens[i]:
+            print(f"Mismatch at {i}: {reference_ids[i]} != {generated_tokens[i]}")
+            break
+    print(f"Match: {reference_ids == generated_tokens}")

nanochat/execution.py

@@ -0,0 +1,349 @@
+"""
+Sandboxed execution utilities for running Python code that comes out of an LLM.
+Adapted from OpenAI HumanEval code:
+https://github.com/openai/human-eval/blob/master/human_eval/execution.py
+
+What is covered:
+- Each execution runs in its own process (can be killed if it hangs or crashes)
+- Execution is limited by a timeout to stop infinite loops
+- Memory limits are enforced by default (256MB)
+- stdout and stderr are captured and returned
+- Code runs in a temporary directory that is deleted afterwards
+- Dangerous functions are disabled (examples: os.system, os.kill, shutil.rmtree, subprocess.Popen)
+
+What is not covered:
+- Not a true security sandbox
+- Network access is not blocked (e.g. sockets could be opened)
+- Python's dynamic features (e.g. ctypes) could bypass restrictions
+- No kernel-level isolation (no seccomp, no containers, no virtualization)
+
+Overall this sandbox is good for evaluation of generated code and protects against
+accidental destructive behavior, but it is not safe against malicious adversarial code.
+"""
+
+import contextlib
+import faulthandler
+import io
+import multiprocessing
+import os
+import platform
+import signal
+import tempfile
+from dataclasses import dataclass
+from typing import Optional
+
+# -----------------------------------------------------------------------------
+
+@dataclass
+class ExecutionResult:
+    """Result of executing Python code in a sandbox."""
+    success: bool
+    stdout: str
+    stderr: str
+    error: Optional[str] = None
+    timeout: bool = False
+    memory_exceeded: bool = False
+
+    def __repr__(self):
+        parts = []
+        parts.append(f"ExecutionResult(success={self.success}")
+        if self.timeout:
+            parts.append(", timeout=True")
+        if self.memory_exceeded:
+            parts.append(", memory_exceeded=True")
+        if self.error:
+            parts.append(f", error={self.error!r}")
+        if self.stdout:
+            parts.append(f", stdout={self.stdout!r}")
+        if self.stderr:
+            parts.append(f", stderr={self.stderr!r}")
+        parts.append(")")
+        return "".join(parts)
+
+
+@contextlib.contextmanager
+def time_limit(seconds: float):
+    def signal_handler(signum, frame):
+        raise TimeoutException("Timed out!")
+
+    signal.setitimer(signal.ITIMER_REAL, seconds)
+    signal.signal(signal.SIGALRM, signal_handler)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+
+
+@contextlib.contextmanager
+def capture_io():
+    """Capture stdout and stderr, and disable stdin."""
+    stdout_capture = io.StringIO()
+    stderr_capture = io.StringIO()
+    stdin_block = WriteOnlyStringIO()
+    with contextlib.redirect_stdout(stdout_capture):
+        with contextlib.redirect_stderr(stderr_capture):
+            with redirect_stdin(stdin_block):
+                yield stdout_capture, stderr_capture
+
+
+@contextlib.contextmanager
+def create_tempdir():
+    with tempfile.TemporaryDirectory() as dirname:
+        with chdir(dirname):
+            yield dirname
+
+
+class TimeoutException(Exception):
+    pass
+
+
+class WriteOnlyStringIO(io.StringIO):
+    """StringIO that throws an exception when it's read from"""
+
+    def read(self, *args, **kwargs):
+        raise IOError
+
+    def readline(self, *args, **kwargs):
+        raise IOError
+
+    def readlines(self, *args, **kwargs):
+        raise IOError
+
+    def readable(self, *args, **kwargs):
+        """Returns True if the IO object can be read."""
+        return False
+
+
+class redirect_stdin(contextlib._RedirectStream):  # type: ignore
+    _stream = "stdin"
+
+
+@contextlib.contextmanager
+def chdir(root):
+    if root == ".":
+        yield
+        return
+    cwd = os.getcwd()
+    os.chdir(root)
+    try:
+        yield
+    finally:
+        os.chdir(cwd)
+
+
+def reliability_guard(maximum_memory_bytes: Optional[int] = None):
+    """
+    This disables various destructive functions and prevents the generated code
+    from interfering with the test (e.g. fork bomb, killing other processes,
+    removing filesystem files, etc.)
+
+    WARNING
+    This function is NOT a security sandbox. Untrusted code, including, model-
+    generated code, should not be blindly executed outside of one. See the
+    Codex paper for more information about OpenAI's code sandbox, and proceed
+    with caution.
+    """
+
+    if platform.uname().system != "Darwin":
+        # These resource limit calls seem to fail on macOS (Darwin), skip?
+        import resource
+        resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))
+        resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))
+        resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))
+
+    faulthandler.disable()
+
+    import builtins
+
+    builtins.exit = None
+    builtins.quit = None
+
+    import os
+
+    os.environ["OMP_NUM_THREADS"] = "1"
+
+    os.kill = None
+    os.system = None
+    os.putenv = None
+    os.remove = None
+    os.removedirs = None
+    os.rmdir = None
+    os.fchdir = None
+    os.setuid = None
+    os.fork = None
+    os.forkpty = None
+    os.killpg = None
+    os.rename = None
+    os.renames = None
+    os.truncate = None
+    os.replace = None
+    os.unlink = None
+    os.fchmod = None
+    os.fchown = None
+    os.chmod = None
+    os.chown = None
+    os.chroot = None
+    os.fchdir = None
+    os.lchflags = None
+    os.lchmod = None
+    os.lchown = None
+    os.getcwd = None
+    os.chdir = None
+
+    import shutil
+
+    shutil.rmtree = None
+    shutil.move = None
+    shutil.chown = None
+
+    import subprocess
+
+    subprocess.Popen = None  # type: ignore
+
+    __builtins__["help"] = None
+
+    import sys
+
+    sys.modules["ipdb"] = None
+    sys.modules["joblib"] = None
+    sys.modules["resource"] = None
+    sys.modules["psutil"] = None
+    sys.modules["tkinter"] = None
+
+
+def _unsafe_execute(code: str, timeout: float, maximum_memory_bytes: Optional[int], result_dict):
+    """Execute code in a subprocess with safety guards. Results are written to result_dict."""
+    with create_tempdir():
+
+        # These system calls are needed when cleaning up tempdir.
+        import os
+        import shutil
+
+        rmtree = shutil.rmtree
+        rmdir = os.rmdir
+        chdir = os.chdir
+        unlink = os.unlink
+
+        # Disable functionalities that can make destructive changes to the test.
+        reliability_guard(maximum_memory_bytes=maximum_memory_bytes)
+
+        # Default to failure
+        result_dict.update({
+            "success": False,
+            "stdout": "",
+            "stderr": "",
+            "timeout": False,
+            "memory_exceeded": False,
+            "error": None,
+        })
+
+        try:
+            exec_globals = {}
+            with capture_io() as (stdout_capture, stderr_capture):
+                with time_limit(timeout):
+                    # WARNING
+                    # This program exists to execute untrusted model-generated code. Although
+                    # it is highly unlikely that model-generated code will do something overtly
+                    # malicious in response to this test suite, model-generated code may act
+                    # destructively due to a lack of model capability or alignment.
+                    # Users are strongly encouraged to sandbox this evaluation suite so that it
+                    # does not perform destructive actions on their host or network. For more
+                    # information on how OpenAI sandboxes its code, see the accompanying paper.
+                    # Once you have read this disclaimer and taken appropriate precautions,
+                    # uncomment the following line and proceed at your own risk:
+                    exec(code, exec_globals)
+
+            result_dict.update({
+                "success": True,
+                "stdout": stdout_capture.getvalue(),
+                "stderr": stderr_capture.getvalue(),
+            })
+
+        except TimeoutException:
+            result_dict.update({
+                "timeout": True,
+                "error": "Execution timed out",
+            })
+
+        except MemoryError as e:
+            result_dict.update({
+                "memory_exceeded": True,
+                "error": f"Memory limit exceeded: {e}",
+            })
+
+        except BaseException as e:
+            result_dict.update({
+                "error": f"{type(e).__name__}: {e}",
+            })
+
+        # Needed for cleaning up.
+        shutil.rmtree = rmtree
+        os.rmdir = rmdir
+        os.chdir = chdir
+        os.unlink = unlink
+
+
+def execute_code(
+    code: str,
+    timeout: float = 5.0, # 5 seconds default
+    maximum_memory_bytes: Optional[int] = 256 * 1024 * 1024, # 256MB default
+) -> ExecutionResult:
+    """
+    Execute Python code in a sandboxed environment.
+
+    Args:
+        code: Python code to execute as a string
+        timeout: Maximum execution time in seconds (default: 5.0)
+        maximum_memory_bytes: Memory limit in bytes (default: 256MB, None to disable)
+
+    Returns:
+        ExecutionResult with success status, stdout/stderr, and error information
+
+    Example:
+        >>> result = execute_code("print('hello world')")
+        >>> result.success
+        True
+        >>> result.stdout
+        'hello world\\n'
+    """
+
+    manager = multiprocessing.Manager()
+    result_dict = manager.dict()
+
+    p = multiprocessing.Process(
+        target=_unsafe_execute,
+        args=(code, timeout, maximum_memory_bytes, result_dict)
+    )
+    p.start()
+    p.join(timeout=timeout + 1)
+
+    if p.is_alive():
+        p.kill()
+        return ExecutionResult(
+            success=False,
+            stdout="",
+            stderr="",
+            error="Execution timed out (process killed)",
+            timeout=True,
+            memory_exceeded=False,
+        )
+
+    if not result_dict:
+        return ExecutionResult(
+            success=False,
+            stdout="",
+            stderr="",
+            error="Execution failed (no result returned)",
+            timeout=True,
+            memory_exceeded=False,
+        )
+
+    return ExecutionResult(
+        success=result_dict["success"],
+        stdout=result_dict["stdout"],
+        stderr=result_dict["stderr"],
+        error=result_dict["error"],
+        timeout=result_dict["timeout"],
+        memory_exceeded=result_dict["memory_exceeded"],
+    )
+

nanochat/gpt.py

@@ -0,0 +1,307 @@
+"""
+GPT model (rewrite, a lot simpler)
+Notable features:
+- rotary embeddings (and no positional embeddings)
+- QK norm
+- untied weights for token embedding and lm_head
+- relu^2 activation in MLP
+- norm after token embedding
+- no learnable params in rmsnorm
+- no bias in linear layers
+- Multi-Query Attention (MQA) support for more efficient inference
+"""
+
+import math
+from functools import partial
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nanochat.common import get_dist_info, print0
+from nanochat.muon import Muon, DistMuon
+from nanochat.adamw import DistAdamW
+
+@dataclass
+class GPTConfig:
+    sequence_len: int = 1024
+    vocab_size: int = 50304
+    n_layer: int = 12
+    n_head: int = 6 # number of query heads
+    n_kv_head: int = 6 # number of key/value heads (MQA)
+    n_embd: int = 768
+
+
+def norm(x):
+    # Purely functional rmsnorm with no learnable params
+    return F.rms_norm(x, (x.size(-1),))
+
+
+def apply_rotary_emb(x, cos, sin):
+    assert x.ndim == 4  # multihead attention
+    d = x.shape[3] // 2
+    x1, x2 = x[..., :d], x[..., d:] # split up last time into two halves
+    y1 = x1 * cos + x2 * sin # rotate pairs of dims
+    y2 = x1 * (-sin) + x2 * cos
+    out = torch.cat([y1, y2], 3) # re-assemble
+    out = out.to(x.dtype) # ensure input/output dtypes match
+    return out
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.n_embd = config.n_embd
+        self.head_dim = self.n_embd // self.n_head
+        assert self.n_embd % self.n_head == 0
+        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
+        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
+
+    def forward(self, x, cos_sin, kv_cache):
+        B, T, C = x.size()
+
+        # Project the input to get queries, keys, and values
+        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
+        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
+        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
+
+        # Apply Rotary Embeddings to queries and keys to get relative positional encoding
+        cos, sin = cos_sin
+        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) # QK rotary embedding
+        q, k = norm(q), norm(k) # QK norm
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) # make head be batch dim, i.e. (B, T, H, D) -> (B, H, T, D)
+
+        # Apply KV cache: insert current k,v into cache, get the full view so far
+        if kv_cache is not None:
+            k, v = kv_cache.insert_kv(self.layer_idx, k, v)
+        Tq = q.size(2) # number of queries in this forward pass
+        Tk = k.size(2) # number of keys/values in total (in the cache + current forward pass)
+
+        # Attention: queries attend to keys/values autoregressively. A few cases to handle:
+        enable_gqa = self.n_head != self.n_kv_head # Group Query Attention (GQA): duplicate key/value heads to match query heads if desired
+        if kv_cache is None or Tq == Tk:
+            # During training (no KV cache), attend as usual with causal attention
+            # And even if there is KV cache, we can still use this simple version when Tq == Tk
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=enable_gqa)
+        elif Tq == 1:
+            # During inference but with a single query in this forward pass:
+            # The query has to attend to all the keys/values in the cache
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=enable_gqa)
+        else:
+            # During inference AND we have a chunk of queries in this forward pass:
+            # First, each query attends to all the cached keys/values (i.e. full prefix)
+            attn_mask = torch.zeros((Tq, Tk), dtype=torch.bool, device=q.device) # True = keep, False = mask
+            prefix_len = Tk - Tq
+            if prefix_len > 0: # can't be negative but could be zero
+                attn_mask[:, :prefix_len] = True
+            # Then, causal attention within this chunk
+            attn_mask[:, prefix_len:] = torch.tril(torch.ones((Tq, Tq), dtype=torch.bool, device=q.device))
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, enable_gqa=enable_gqa)
+
+        # Re-assemble the heads side by side and project back to residual stream
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
+        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.attn = CausalSelfAttention(config, layer_idx)
+        self.mlp = MLP(config)
+
+    def forward(self, x, cos_sin, kv_cache):
+        x = x + self.attn(norm(x), cos_sin, kv_cache)
+        x = x + self.mlp(norm(x))
+        return x
+
+
+class GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.transformer = nn.ModuleDict({
+            "wte": nn.Embedding(config.vocab_size, config.n_embd),
+            "h": nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.n_layer)]),
+        })
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # To support meta device initialization, we init the rotary embeddings here, but it's fake
+        # As for rotary_seq_len, these rotary embeddings are pretty small/cheap in memory,
+        # so let's just over-compute them, but assert fail if we ever reach that amount.
+        # In the future we can dynamically grow the cache, for now it's fine.
+        self.rotary_seq_len = config.sequence_len * 10 # 10X over-compute should be enough, TODO make nicer?
+        head_dim = config.n_embd // config.n_head
+        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+        self.register_buffer("cos", cos, persistent=False) # persistent=False means it's not saved to the checkpoint
+        self.register_buffer("sin", sin, persistent=False)
+
+    def init_weights(self):
+        self.apply(self._init_weights)
+        # zero out classifier weights
+        torch.nn.init.zeros_(self.lm_head.weight)
+        # zero out c_proj weights in all blocks
+        for block in self.transformer.h:
+            torch.nn.init.zeros_(block.mlp.c_proj.weight)
+            torch.nn.init.zeros_(block.attn.c_proj.weight)
+        # init the rotary embeddings
+        head_dim = self.config.n_embd // self.config.n_head
+        cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
+        self.cos, self.sin = cos, sin
+        # Cast the embeddings from fp32 to bf16: optim can tolerate it and it saves memory: both in the model and the activations
+        if self.transformer.wte.weight.device.type == "cuda":
+            self.transformer.wte.to(dtype=torch.bfloat16)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            # https://arxiv.org/pdf/2310.17813
+            fan_out = module.weight.size(0)
+            fan_in = module.weight.size(1)
+            std = 1.0 / math.sqrt(fan_in) * min(1.0, math.sqrt(fan_out / fan_in))
+            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=1.0)
+
+    # TODO: bump base theta more, e.g. 100K is more common more recently
+    def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None):
+        # autodetect the device from model embeddings
+        if device is None:
+            device = self.transformer.wte.weight.device
+        # stride the channels
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        # stride the time steps
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        # calculate the rotation frequencies at each (time, channel) pair
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16() # keep them in bfloat16
+        cos, sin = cos[None, :, None, :], sin[None, :, None, :] # add batch and head dims for later broadcasting
+        return cos, sin
+
+    def get_device(self):
+        return self.transformer.wte.weight.device
+
+    def estimate_flops(self):
+        """ Return the estimated FLOPs per token for the model. Ref: https://arxiv.org/abs/2204.02311 """
+        nparams = sum(p.numel() for p in self.parameters())
+        nparams_embedding = self.transformer.wte.weight.numel()
+        l, h, q, t = self.config.n_layer, self.config.n_head, self.config.n_embd // self.config.n_head, self.config.sequence_len
+        num_flops_per_token = 6 * (nparams - nparams_embedding) + 12 * l * h * q * t
+        return num_flops_per_token
+
+    def setup_optimizers(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0):
+        model_dim = self.config.n_embd
+        ddp, rank, local_rank, world_size = get_dist_info()
+        # Separate out all parameters into 3 groups (matrix, embedding, lm_head)
+        matrix_params = list(self.transformer.h.parameters())
+        embedding_params = list(self.transformer.wte.parameters())
+        lm_head_params = list(self.lm_head.parameters())
+        assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params)
+        # Create the AdamW optimizer for the embedding and lm_head
+        # Scale the LR for the AdamW parameters by ∝1/√dmodel (having tuned the LRs for 768 dim model)
+        dmodel_lr_scale = (model_dim / 768) ** -0.5
+        if rank == 0:
+            print(f"Scaling the LR for the AdamW parameters ∝1/√({model_dim}/768) = {dmodel_lr_scale:.6f}")
+        adam_groups = [
+            dict(params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale),
+            dict(params=embedding_params, lr=embedding_lr * dmodel_lr_scale),
+        ]
+        adamw_kwargs = dict(betas=(0.8, 0.95), eps=1e-10, weight_decay=weight_decay)
+        AdamWFactory = DistAdamW if ddp else partial(torch.optim.AdamW, fused=True)
+        adamw_optimizer = AdamWFactory(adam_groups, **adamw_kwargs)
+        # Create the Muon optimizer for the linear layers
+        muon_kwargs = dict(lr=matrix_lr, momentum=0.95)
+        MuonFactory = DistMuon if ddp else Muon
+        muon_optimizer = MuonFactory(matrix_params, **muon_kwargs)
+        # Combine them the two optimizers into one list
+        optimizers = [adamw_optimizer, muon_optimizer]
+        for opt in optimizers:
+            for group in opt.param_groups:
+                group["initial_lr"] = group["lr"]
+        return optimizers
+
+    def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
+        B, T = idx.size()
+
+        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim))
+        assert T <= self.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {self.cos.size(1)}"
+        assert idx.device == self.cos.device, f"Rotary embeddings and idx are on different devices: {idx.device} != {self.cos.device}"
+        assert self.cos.dtype == torch.bfloat16, "Rotary embeddings must be in bfloat16"
+        # if kv cache exists, we need to offset the rotary embeddings to the current position in the cache
+        T0 = 0 if kv_cache is None else kv_cache.get_pos()
+        cos_sin = self.cos[:, T0:T0+T], self.sin[:, T0:T0+T] # truncate cache to current sequence length
+
+        # Forward the trunk of the Transformer
+        x = self.transformer.wte(idx)
+        x = norm(x)
+        for block in self.transformer.h:
+            x = block(x, cos_sin, kv_cache)
+        x = norm(x)
+
+        # Forward the lm_head (compute logits)
+        softcap = 15
+        if targets is not None:
+            # training mode: compute and return the loss
+            # TODO: experiment with Liger Kernels / chunked cross-entropy etc.
+            logits = self.lm_head(x)
+            logits = softcap * torch.tanh(logits / softcap) # logits softcap
+            logits = logits.float() # use tf32/fp32 for logits
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction)
+            return loss
+        else:
+            # inference mode: compute and return the logits
+            logits = self.lm_head(x)
+            logits = softcap * torch.tanh(logits / softcap) # logits softcap
+            return logits
+
+    @torch.inference_mode()
+    def generate(self, tokens, max_tokens, temperature=1.0, top_k=None, seed=42):
+        """
+        Naive autoregressive streaming inference.
+        To make it super simple, let's assume:
+        - batch size is 1
+        - ids and the yielded tokens are simple Python lists and ints
+        """
+        assert isinstance(tokens, list)
+        device = self.get_device()
+        rng = None
+        if temperature > 0:
+            rng = torch.Generator(device=device)
+            rng.manual_seed(seed)
+        ids = torch.tensor([tokens], dtype=torch.long, device=device) # add batch dim
+        for _ in range(max_tokens):
+            logits = self.forward(ids) # (B, T, vocab_size)
+            logits = logits[:, -1, :] # (B, vocab_size)
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            if temperature > 0:
+                logits = logits / temperature
+                probs = F.softmax(logits, dim=-1)
+                next_ids = torch.multinomial(probs, num_samples=1, generator=rng)
+            else:
+                next_ids = torch.argmax(logits, dim=-1, keepdim=True)
+            ids = torch.cat((ids, next_ids), dim=1)
+            token = next_ids.item()
+            yield token

nanochat/loss_eval.py

@@ -0,0 +1,65 @@
+"""
+A number of functions that help with evaluating a base model.
+"""
+import math
+import torch
+import torch.distributed as dist
+
+@torch.no_grad()
+def evaluate_bpb(model, batches, steps, token_bytes):
+    """
+    Instead of the naive 'mean loss', this function returns the bits per byte (bpb),
+    which is a tokenization vocab size-indepedent metric, meaning you are still comparing
+    apples:apples if you change the vocab size. The way this works is that instead of just
+    calculating the average loss as usual, you calculate the sum loss, and indepependently
+    also the sum bytes (of all the target tokens), and divide. This normalizes the loss by
+    the number of bytes that the target tokens represent.
+
+    The added complexity is so that:
+    1) All "normal" tokens are normalized by the length of the token in bytes
+    2) No special tokens (e.g. <|bos|>) are included in the metric - they are masked out.
+    3) No actively masked tokens (using ignore_index of e.g. -1) are included in the metric.
+
+    In addition to evaluate_loss, we need the token_bytes tensor:
+    It is a 1D tensor of shape (vocab_size,), indicating the number of bytes for
+    each token id, or 0 if the token is to not be counted (e.g. special tokens).
+    """
+    # record the losses
+    total_nats = torch.tensor(0.0, dtype=torch.float32, device=model.get_device())
+    total_bytes = torch.tensor(0, dtype=torch.int64, device=model.get_device())
+    batch_iter = iter(batches)
+    for _ in range(steps):
+        x, y = next(batch_iter)
+        loss2d = model(x, y, loss_reduction='none') # (B, T)
+        loss2d = loss2d.view(-1) # flatten
+        y = y.view(-1) # flatten
+        if (y.int() < 0).any(): # mps does not currently have kernel for < 0 for int64, only int32
+            # slightly more complex code path if some target tokens are ignore_index (e.g. -1)
+            # any target token < 0 is to be ignored: do NOT index token_bytes with negatives
+            valid = y >= 0
+            y_safe = torch.where(valid, y, torch.zeros_like(y))
+            # map valid targets to their byte length; ignored targets contribute 0 bytes
+            num_bytes2d = torch.where(
+                valid,
+                token_bytes[y_safe],
+                torch.zeros_like(y, dtype=token_bytes.dtype)
+            )
+            total_nats += (loss2d * (num_bytes2d > 0)).sum()
+            total_bytes += num_bytes2d.sum()
+        else:
+            # fast path: no ignored targets, safe to index directly
+            num_bytes2d = token_bytes[y]
+            total_nats += (loss2d * (num_bytes2d > 0)).sum()
+            total_bytes += num_bytes2d.sum()
+    # sum reduce across all ranks
+    world_size = dist.get_world_size() if dist.is_initialized() else 1
+    if world_size > 1:
+        dist.all_reduce(total_nats, op=dist.ReduceOp.SUM)
+        dist.all_reduce(total_bytes, op=dist.ReduceOp.SUM)
+    # move both to cpu, calculate bpb and return
+    total_nats = total_nats.item()
+    total_bytes = total_bytes.item()
+    if total_bytes == 0:
+        return float('inf')
+    bpb = total_nats / (math.log(2) * total_bytes)
+    return bpb

nanochat/muon.py

@@ -0,0 +1,187 @@
+"""
+Muon optimizer from Keller et al.
+Also a lot of borrowing of ideas from modded-nanogpt.
+"""
+import torch
+from torch import Tensor
+import torch.distributed as dist
+
+@torch.compile
+def zeropower_via_newtonschulz5(G: Tensor, steps: int) -> Tensor:
+    """
+    Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
+    quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
+    of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
+    zero even beyond the point where the iteration no longer converges all the way to one everywhere
+    on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
+    where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+    """
+    assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
+    a, b, c = (3.4445, -4.7750,  2.0315)
+    X = G.bfloat16()
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
+    # Perform the NS iterations
+    for _ in range(steps):
+        A = X @ X.mT
+        B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
+
+    if G.size(-2) > G.size(-1):
+        X = X.mT
+    return X
+
+class Muon(torch.optim.Optimizer):
+    """
+    Muon - MomentUm Orthogonalized by Newton-schulz
+
+    https://kellerjordan.github.io/posts/muon/
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
+    processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
+    matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has
+    the advantage that it can be stably run in bfloat16 on the GPU.
+
+    Some warnings:
+    - This optimizer should not be used for the embedding layer, the final fully connected layer,
+    or any {0,1}-D parameters; those should all be optimized by a standard method (e.g., AdamW).
+    - To use it with 4D convolutional filters, it works well to just flatten their last 3 dimensions.
+
+    Arguments:
+        lr: The learning rate used by the internal SGD.
+        momentum: The momentum used by the internal SGD.
+        nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended)
+        ns_steps: The number of Newton-Schulz iteration steps to use.
+    """
+    def __init__(self, params, lr=0.02, momentum=0.95, nesterov=True, ns_steps=5):
+        defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
+        params: list[Tensor] = [*params]
+        param_groups = []
+        for size in {p.numel() for p in params}:
+            group = dict(params=[p for p in params if p.numel() == size])
+            param_groups.append(group)
+        super().__init__(param_groups, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        for group in self.param_groups:
+            params: list[Tensor] = group["params"]
+            for p in params:
+                g = p.grad
+                assert g is not None
+                state = self.state[p]
+                if "momentum_buffer" not in state:
+                    state["momentum_buffer"] = torch.zeros_like(g)
+                buf: Tensor = state["momentum_buffer"]
+                buf.lerp_(g, 1 - group["momentum"])
+                g = g.lerp_(buf, group["momentum"]) if group["nesterov"] else buf
+                g = zeropower_via_newtonschulz5(g, steps=group["ns_steps"])
+                p.add_(g, alpha=-group["lr"] * max(1, p.size(-2) / p.size(-1))**0.5)
+
+
+class DistMuon(torch.optim.Optimizer):
+    """
+    Muon: SGD-momentum + (optional) Nesterov, then orthogonalize the 2D update via Newton–Schulz,
+    finally apply aspect-ratio scaled step. Performs its own distributed synchronization:
+      - reduce_scatter(AVG) for gradient averaging
+      - all_gather to replicate updated weights
+
+    Notes:
+      * Designed for 2D parameters (e.g., linear/conv kernels reshaped to 2D). Do not use for 0D/1D
+        params like embeddings or scalars.
+      * Momentum buffers are maintained only on the 'owner' rank for each parameter (rank chosen
+        by block-cyclic assignment below). If you checkpoint optimizer state on a single rank,
+        consolidate states beforehand.
+
+    Args:
+        params: iterable of Tensors
+        lr: learning rate
+        momentum: momentum coefficient in [0,1)
+        nesterov: if True, Nesterov-style update (g <- lerp(g, buf, momentum)); else use buf
+        ns_steps: number of Newton–Schulz iterations for the orthogonalization
+    """
+    def __init__(self, params, lr: float = 0.02, momentum: float = 0.95,
+                 nesterov: bool = True, ns_steps: int = 5):
+        defaults = dict(lr=lr, momentum=momentum, nesterov=nesterov, ns_steps=ns_steps)
+        params = list(params)
+        assert all(p.ndim == 2 for p in params), "Muon expects 2D parameters only"
+        rank = dist.get_rank()
+        # Group all parameters by their shape
+        shapes = sorted({p.shape for p in params}) # sort to ensure consistent / deterministic ordering
+        param_groups = []
+        for shape in shapes:
+            group_params = [p for p in params if p.shape == shape]
+            device, dtype = group_params[0].device, group_params[0].dtype
+            assert all(p.device == device for p in group_params)
+            assert all(p.dtype == dtype for p in group_params)
+            if rank == 0:
+                print(f"Muon: Grouping {len(group_params)} params of shape {shape}, device {device}, dtype {dtype}")
+            param_groups.append(dict(params=group_params, zero_buffer=torch.zeros_like(group_params[0])))
+        super().__init__(param_groups, defaults)
+
+    @torch.no_grad()
+    def step(self):
+        rank = dist.get_rank()
+        world_size = dist.get_world_size()
+
+        # Ensure all grads exist
+        assert all(p.grad is not None for group in self.param_groups for p in group["params"]), "All params must have grads"
+
+        # Kick off all the reduce scatter operations to average up the gradients across all ranks
+        all_reduce_futures = []
+        for group in self.param_groups:
+            params = group["params"]
+            zero_buffer = group["zero_buffer"]
+            # Go through params in groups of world_size.
+            for base_i in range(0, len(params), world_size):
+                # The compute owner of each param is rank i % world_size
+                owner_idx = base_i + rank
+                # each rank stacks up its chunk of world_size params into a list
+                rs_input = [p.grad for p in params[base_i:base_i + world_size]]
+                # pad rs_input with the zero buffer to complete the group
+                rs_input.extend([zero_buffer] * (world_size - len(rs_input)))
+                # the output buffer gets strided across the group based on the rank
+                rs_output = params[owner_idx].grad if owner_idx < len(params) else torch.empty_like(zero_buffer)
+                # reduce scatter the gradients within this group of world_size params
+                work = dist.reduce_scatter(rs_output, rs_input, op=dist.ReduceOp.AVG, async_op=True).get_future()
+                all_reduce_futures.append(work)
+
+        # Now each rank computes the update and gathers
+        future_idx = 0
+        all_gather_futures = []
+        for group in self.param_groups:
+            params = group["params"]
+            zero_buffer = group["zero_buffer"]
+            # Go through params in groups of world_size.
+            for base_i in range(0, len(params), world_size):
+                # The compute owner of each param is rank i % world_size
+                owner_idx = base_i + rank # calculate the index of the param that this rank owns
+                # Wait for the reduce scatter to complete
+                all_reduce_futures[future_idx].wait() # possibly later we could use wait_any polling instead
+                future_idx += 1
+                # Owner computes the Muon update, result is in its param
+                if owner_idx < len(params):
+                    p = params[owner_idx]
+                    g = p.grad  # now averaged across ranks
+                    state = self.state[p]
+                    if "momentum_buffer" not in state:
+                        state["momentum_buffer"] = torch.zeros_like(g)
+                    buf: Tensor = state["momentum_buffer"]
+                    buf.lerp_(g, 1.0 - group["momentum"])
+                    g = g.lerp_(buf, group["momentum"]) if group["nesterov"] else buf
+                    g = zeropower_via_newtonschulz5(g, steps=group["ns_steps"])
+                    scale = (max(1.0, p.size(-2) / p.size(-1)) ** 0.5)
+                    p.add_(g, alpha=-group["lr"] * scale)
+                # Replicate updated parameters to all ranks
+                ag_input = params[owner_idx] if owner_idx < len(params) else zero_buffer
+                ag_output = params[base_i:base_i + world_size]
+                ag_output.extend([torch.empty_like(zero_buffer) for _ in range(world_size - len(ag_output))]) # pad
+                work = dist.all_gather(ag_output, ag_input, async_op=True).get_future()
+                all_gather_futures.append(work)
+
+        # Wait for all work to finish
+        torch.futures.collect_all(all_gather_futures).wait()

nanochat/report.py

@@ -0,0 +1,408 @@
+"""
+Utilities for generating training report cards. More messy code than usual, will fix.
+"""
+
+import os
+import re
+import shutil
+import subprocess
+import socket
+import datetime
+import platform
+import psutil
+import torch
+
+def run_command(cmd):
+    """Run a shell command and return output, or None if it fails."""
+    try:
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=5)
+        if result.returncode == 0:
+            return result.stdout.strip()
+        return None
+    except:
+        return None
+
+def get_git_info():
+    """Get current git commit, branch, and dirty status."""
+    info = {}
+    info['commit'] = run_command("git rev-parse --short HEAD") or "unknown"
+    info['branch'] = run_command("git rev-parse --abbrev-ref HEAD") or "unknown"
+
+    # Check if repo is dirty (has uncommitted changes)
+    status = run_command("git status --porcelain")
+    info['dirty'] = bool(status) if status is not None else False
+
+    # Get commit message
+    info['message'] = run_command("git log -1 --pretty=%B") or ""
+    info['message'] = info['message'].split('\n')[0][:80]  # First line, truncated
+
+    return info
+
+def get_gpu_info():
+    """Get GPU information."""
+    if not torch.cuda.is_available():
+        return {"available": False}
+
+    num_devices = torch.cuda.device_count()
+    info = {
+        "available": True,
+        "count": num_devices,
+        "names": [],
+        "memory_gb": []
+    }
+
+    for i in range(num_devices):
+        props = torch.cuda.get_device_properties(i)
+        info["names"].append(props.name)
+        info["memory_gb"].append(props.total_memory / (1024**3))
+
+    # Get CUDA version
+    info["cuda_version"] = torch.version.cuda or "unknown"
+
+    return info
+
+def get_system_info():
+    """Get system information."""
+    info = {}
+
+    # Basic system info
+    info['hostname'] = socket.gethostname()
+    info['platform'] = platform.system()
+    info['python_version'] = platform.python_version()
+    info['torch_version'] = torch.__version__
+
+    # CPU and memory
+    info['cpu_count'] = psutil.cpu_count(logical=False)
+    info['cpu_count_logical'] = psutil.cpu_count(logical=True)
+    info['memory_gb'] = psutil.virtual_memory().total / (1024**3)
+
+    # User and environment
+    info['user'] = os.environ.get('USER', 'unknown')
+    info['nanochat_base_dir'] = os.environ.get('NANOCHAT_BASE_DIR', 'out')
+    info['working_dir'] = os.getcwd()
+
+    return info
+
+def estimate_cost(gpu_info, runtime_hours=None):
+    """Estimate training cost based on GPU type and runtime."""
+
+    # Rough pricing, from Lambda Cloud
+    default_rate = 2.0
+    gpu_hourly_rates = {
+        "H100": 3.00,
+        "A100": 1.79,
+        "V100": 0.55,
+    }
+
+    if not gpu_info.get("available"):
+        return None
+
+    # Try to identify GPU type from name
+    hourly_rate = None
+    gpu_name = gpu_info["names"][0] if gpu_info["names"] else "unknown"
+    for gpu_type, rate in gpu_hourly_rates.items():
+        if gpu_type in gpu_name:
+            hourly_rate = rate * gpu_info["count"]
+            break
+
+    if hourly_rate is None:
+        hourly_rate = default_rate * gpu_info["count"]  # Default estimate
+
+    return {
+        "hourly_rate": hourly_rate,
+        "gpu_type": gpu_name,
+        "estimated_total": hourly_rate * runtime_hours if runtime_hours else None
+    }
+
+def generate_header():
+    """Generate the header for a training report."""
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+
+    git_info = get_git_info()
+    gpu_info = get_gpu_info()
+    sys_info = get_system_info()
+    cost_info = estimate_cost(gpu_info)
+
+    header = f"""# nanochat training report
+
+Generated: {timestamp}
+
+## Environment
+
+### Git Information
+- Branch: {git_info['branch']}
+- Commit: {git_info['commit']} {"(dirty)" if git_info['dirty'] else "(clean)"}
+- Message: {git_info['message']}
+
+### Hardware
+- Platform: {sys_info['platform']}
+- CPUs: {sys_info['cpu_count']} cores ({sys_info['cpu_count_logical']} logical)
+- Memory: {sys_info['memory_gb']:.1f} GB
+"""
+
+    if gpu_info.get("available"):
+        gpu_names = ", ".join(set(gpu_info["names"]))
+        total_vram = sum(gpu_info["memory_gb"])
+        header += f"""- GPUs: {gpu_info['count']}x {gpu_names}
+- GPU Memory: {total_vram:.1f} GB total
+- CUDA Version: {gpu_info['cuda_version']}
+"""
+    else:
+        header += "- GPUs: None available\n"
+
+    if cost_info and cost_info["hourly_rate"] > 0:
+        header += f"""- Hourly Rate: ${cost_info['hourly_rate']:.2f}/hour\n"""
+
+    header += f"""
+### Software
+- Python: {sys_info['python_version']}
+- PyTorch: {sys_info['torch_version']}
+
+"""
+
+    # bloat metrics: package all of the source code and assess its weight
+    packaged = run_command('files-to-prompt . -e py -e md -e rs -e html -e toml -e sh --ignore "*target*" --cxml')
+    num_chars = len(packaged)
+    num_lines = len(packaged.split('\n'))
+    num_files = len([x for x in packaged.split('\n') if x.startswith('<source>')])
+    num_tokens = num_chars // 4 # assume approximately 4 chars per token
+
+    # count dependencies via uv.lock
+    uv_lock_lines = 0
+    if os.path.exists('uv.lock'):
+        with open('uv.lock', 'r', encoding='utf-8') as f:
+            uv_lock_lines = len(f.readlines())
+
+    header += f"""
+### Bloat
+- Characters: {num_chars:,}
+- Lines: {num_lines:,}
+- Files: {num_files:,}
+- Tokens (approx): {num_tokens:,}
+- Dependencies (uv.lock lines): {uv_lock_lines:,}
+
+"""
+    return header
+
+# -----------------------------------------------------------------------------
+
+def slugify(text):
+    """Slugify a text string."""
+    return text.lower().replace(" ", "-")
+
+# the expected files and their order
+EXPECTED_FILES = [
+    "tokenizer-training.md",
+    "tokenizer-evaluation.md",
+    "base-model-training.md",
+    "base-model-loss.md",
+    "base-model-evaluation.md",
+    "midtraining.md",
+    "chat-evaluation-mid.md",
+    "chat-sft.md",
+    "chat-evaluation-sft.md",
+    "chat-rl.md",
+    "chat-evaluation-rl.md",
+]
+# the metrics we're currently interested in
+chat_metrics = ["ARC-Easy", "ARC-Challenge", "MMLU", "GSM8K", "HumanEval", "ChatCORE"]
+
+def extract(section, keys):
+    """simple def to extract a single key from a section"""
+    if not isinstance(keys, list):
+        keys = [keys] # convenience
+    out = {}
+    for line in section.split("\n"):
+        for key in keys:
+            if key in line:
+                out[key] = line.split(":")[1].strip()
+    return out
+
+def extract_timestamp(content, prefix):
+    """Extract timestamp from content with given prefix."""
+    for line in content.split('\n'):
+        if line.startswith(prefix):
+            time_str = line.split(":", 1)[1].strip()
+            try:
+                return datetime.datetime.strptime(time_str, "%Y-%m-%d %H:%M:%S")
+            except:
+                pass
+    return None
+
+class Report:
+    """Maintains a bunch of logs, generates a final markdown report."""
+
+    def __init__(self, report_dir):
+        os.makedirs(report_dir, exist_ok=True)
+        self.report_dir = report_dir
+
+    def log(self, section, data):
+        """Log a section of data to the report."""
+        slug = slugify(section)
+        file_name = f"{slug}.md"
+        file_path = os.path.join(self.report_dir, file_name)
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(f"## {section}\n")
+            f.write(f"timestamp: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
+            for item in data:
+                if not item:
+                    # skip falsy values like None or empty dict etc.
+                    continue
+                if isinstance(item, str):
+                    # directly write the string
+                    f.write(item)
+                else:
+                    # render a dict
+                    for k, v in item.items():
+                        if isinstance(v, float):
+                            vstr = f"{v:.4f}"
+                        elif isinstance(v, int) and v >= 10000:
+                            vstr = f"{v:,.0f}"
+                        else:
+                            vstr = str(v)
+                        f.write(f"- {k}: {vstr}\n")
+            f.write("\n")
+        return file_path
+
+    def generate(self):
+        """Generate the final report."""
+        report_dir = self.report_dir
+        report_file = os.path.join(report_dir, "report.md")
+        print(f"Generating report to {report_file}")
+        final_metrics = {} # the most important final metrics we'll add as table at the end
+        start_time = None
+        end_time = None
+        with open(report_file, "w", encoding="utf-8") as out_file:
+            # write the header first
+            header_file = os.path.join(report_dir, "header.md")
+            if os.path.exists(header_file):
+                with open(header_file, "r", encoding="utf-8") as f:
+                    header_content = f.read()
+                    out_file.write(header_content)
+                    start_time = extract_timestamp(header_content, "Run started:")
+                    # capture bloat data for summary later (the stuff after Bloat header and until \n\n)
+                    bloat_data = re.search(r"### Bloat\n(.*?)\n\n", header_content, re.DOTALL)
+                    bloat_data = bloat_data.group(1) if bloat_data else ""
+            else:
+                start_time = None # will cause us to not write the total wall clock time
+                bloat_data = "[bloat data missing]"
+                print(f"Warning: {header_file} does not exist. Did you forget to run `nanochat reset`?")
+            # process all the individual sections
+            for file_name in EXPECTED_FILES:
+                section_file = os.path.join(report_dir, file_name)
+                if not os.path.exists(section_file):
+                    print(f"Warning: {section_file} does not exist, skipping")
+                    continue
+                with open(section_file, "r", encoding="utf-8") as in_file:
+                    section = in_file.read()
+                # Extract timestamp from this section (the last section's timestamp will "stick" as end_time)
+                if "rl" not in file_name:
+                    # Skip RL sections for end_time calculation because RL is experimental
+                    end_time = extract_timestamp(section, "timestamp:")
+                # extract the most important metrics from the sections
+                if file_name == "base-model-evaluation.md":
+                    final_metrics["base"] = extract(section, "CORE")
+                if file_name == "chat-evaluation-mid.md":
+                    final_metrics["mid"] = extract(section, chat_metrics)
+                if file_name == "chat-evaluation-sft.md":
+                    final_metrics["sft"] = extract(section, chat_metrics)
+                if file_name == "chat-evaluation-rl.md":
+                    final_metrics["rl"] = extract(section, "GSM8K") # RL only evals GSM8K
+                # append this section of the report
+                out_file.write(section)
+                out_file.write("\n")
+            # add the final metrics table
+            out_file.write("## Summary\n\n")
+            # Copy over the bloat metrics from the header
+            out_file.write(bloat_data)
+            out_file.write("\n\n")
+            # Collect all unique metric names
+            all_metrics = set()
+            for stage_metrics in final_metrics.values():
+                all_metrics.update(stage_metrics.keys())
+            # Custom ordering: CORE first, ChatCORE last, rest in middle
+            all_metrics = sorted(all_metrics, key=lambda x: (x != "CORE", x == "ChatCORE", x))
+            # Fixed column widths
+            stages = ["base", "mid", "sft", "rl"]
+            metric_width = 15
+            value_width = 8
+            # Write table header
+            header = f"| {'Metric'.ljust(metric_width)} |"
+            for stage in stages:
+                header += f" {stage.upper().ljust(value_width)} |"
+            out_file.write(header + "\n")
+            # Write separator
+            separator = f"|{'-' * (metric_width + 2)}|"
+            for stage in stages:
+                separator += f"{'-' * (value_width + 2)}|"
+            out_file.write(separator + "\n")
+            # Write table rows
+            for metric in all_metrics:
+                row = f"| {metric.ljust(metric_width)} |"
+                for stage in stages:
+                    value = final_metrics.get(stage, {}).get(metric, "-")
+                    row += f" {str(value).ljust(value_width)} |"
+                out_file.write(row + "\n")
+            out_file.write("\n")
+            # Calculate and write total wall clock time
+            if start_time and end_time:
+                duration = end_time - start_time
+                total_seconds = int(duration.total_seconds())
+                hours = total_seconds // 3600
+                minutes = (total_seconds % 3600) // 60
+                out_file.write(f"Total wall clock time: {hours}h{minutes}m\n")
+            else:
+                out_file.write("Total wall clock time: unknown\n")
+        # also cp the report.md file to current directory
+        print(f"Copying report.md to current directory for convenience")
+        shutil.copy(report_file, "report.md")
+        return report_file
+
+    def reset(self):
+        """Reset the report."""
+        # Remove section files
+        for file_name in EXPECTED_FILES:
+            file_path = os.path.join(self.report_dir, file_name)
+            if os.path.exists(file_path):
+                os.remove(file_path)
+        # Remove report.md if it exists
+        report_file = os.path.join(self.report_dir, "report.md")
+        if os.path.exists(report_file):
+            os.remove(report_file)
+        # Generate and write the header section with start timestamp
+        header_file = os.path.join(self.report_dir, "header.md")
+        header = generate_header()
+        start_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        with open(header_file, "w", encoding="utf-8") as f:
+            f.write(header)
+            f.write(f"Run started: {start_time}\n\n---\n\n")
+        print(f"Reset report and wrote header to {header_file}")
+
+# -----------------------------------------------------------------------------
+# nanochat-specific convenience functions
+
+class DummyReport:
+    def log(self, *args, **kwargs):
+        pass
+    def reset(self, *args, **kwargs):
+        pass
+
+def get_report():
+    # just for convenience, only rank 0 logs to report
+    from nanochat.common import get_base_dir, get_dist_info
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
+    if ddp_rank == 0:
+        report_dir = os.path.join(get_base_dir(), "report")
+        return Report(report_dir)
+    else:
+        return DummyReport()
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Generate or reset nanochat training reports.")
+    parser.add_argument("command", nargs="?", default="generate", choices=["generate", "reset"], help="Operation to perform (default: generate)")
+    args = parser.parse_args()
+    if args.command == "generate":
+        get_report().generate()
+    elif args.command == "reset":
+        get_report().reset()

nanochat/tokenizer.py

@@ -0,0 +1,398 @@
+"""
+BPE Tokenizer in the style of GPT-4.
+
+Two implementations are available:
+1) HuggingFace Tokenizer that can do both training and inference but is really confusing
+2) Our own RustBPE Tokenizer for training and tiktoken for efficient inference
+"""
+
+import os
+import copy
+from functools import lru_cache
+
+SPECIAL_TOKENS = [
+    # every document begins with the Beginning of Sequence (BOS) token that delimits documents
+    "<|bos|>",
+    # tokens below are only used during finetuning to render Conversations into token ids
+    "<|user_start|>", # user messages
+    "<|user_end|>",
+    "<|assistant_start|>", # assistant messages
+    "<|assistant_end|>",
+    "<|python_start|>", # assistant invokes python REPL tool
+    "<|python_end|>",
+    "<|output_start|>", # python REPL outputs back to assistant
+    "<|output_end|>",
+]
+
+# NOTE: this split pattern deviates from GPT-4 in that we use \p{N}{1,2} instead of \p{N}{1,3}
+# I did this because I didn't want to "waste" too many tokens on numbers for smaller vocab sizes.
+# I haven't validated that this is actually a good idea, TODO.
+SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""
+
+# -----------------------------------------------------------------------------
+# Generic GPT-4-style tokenizer based on HuggingFace Tokenizer
+from tokenizers import Tokenizer as HFTokenizer
+from tokenizers import pre_tokenizers, decoders, Regex
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+
+class HuggingFaceTokenizer:
+    """Light wrapper around HuggingFace Tokenizer for some utilities"""
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    @classmethod
+    def from_pretrained(cls, hf_path):
+        # init from a HuggingFace pretrained tokenizer (e.g. "gpt2")
+        tokenizer = HFTokenizer.from_pretrained(hf_path)
+        return cls(tokenizer)
+
+    @classmethod
+    def from_directory(cls, tokenizer_dir):
+        # init from a local directory on disk (e.g. "out/tokenizer")
+        tokenizer_path = os.path.join(tokenizer_dir, "tokenizer.json")
+        tokenizer = HFTokenizer.from_file(tokenizer_path)
+        return cls(tokenizer)
+
+    @classmethod
+    def train_from_iterator(cls, text_iterator, vocab_size):
+        # train from an iterator of text
+        # Configure the HuggingFace Tokenizer
+        tokenizer = HFTokenizer(BPE(
+            byte_fallback=True, # needed!
+            unk_token=None,
+            fuse_unk=False,
+        ))
+        # Normalizer: None
+        tokenizer.normalizer = None
+        # Pre-tokenizer: GPT-4 style
+        # the regex pattern used by GPT-4 to split text into groups before BPE
+        # NOTE: The pattern was changed from \p{N}{1,3} to \p{N}{1,2} because I suspect it is harmful to
+        # very small models and smaller vocab sizes, because it is a little bit wasteful in the token space.
+        # (but I haven't validated this! TODO)
+        gpt4_split_regex = Regex(SPLIT_PATTERN) # huggingface demands that you wrap it in Regex!!
+        tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
+            pre_tokenizers.Split(pattern=gpt4_split_regex, behavior="isolated", invert=False),
+            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
+        ])
+        # Decoder: ByteLevel (it pairs together with the ByteLevel pre-tokenizer)
+        tokenizer.decoder = decoders.ByteLevel()
+        # Post-processor: None
+        tokenizer.post_processor = None
+        # Trainer: BPE
+        trainer = BpeTrainer(
+            vocab_size=vocab_size,
+            show_progress=True,
+            min_frequency=0, # no minimum frequency
+            initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+            special_tokens=SPECIAL_TOKENS,
+        )
+        # Kick off the training
+        tokenizer.train_from_iterator(text_iterator, trainer)
+        return cls(tokenizer)
+
+    def get_vocab_size(self):
+        return self.tokenizer.get_vocab_size()
+
+    def get_special_tokens(self):
+        special_tokens_map = self.tokenizer.get_added_tokens_decoder()
+        special_tokens = [w.content for w in special_tokens_map.values()]
+        return special_tokens
+
+    def id_to_token(self, id):
+        return self.tokenizer.id_to_token(id)
+
+    def _encode_one(self, text, prepend=None, append=None):
+        # encode a single string
+        # prepend/append can be either a string of a special token or a token id directly.
+        assert isinstance(text, str)
+        ids = []
+        if prepend is not None:
+            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
+            ids.append(prepend_id)
+        ids.extend(self.tokenizer.encode(text, add_special_tokens=False).ids)
+        if append is not None:
+            append_id = append if isinstance(append, int) else self.encode_special(append)
+            ids.append(append_id)
+        return ids
+
+    def encode_special(self, text):
+        # encode a single special token via exact match
+        return self.tokenizer.token_to_id(text)
+
+    def get_bos_token_id(self):
+        bos = self.encode_special("<|bos|>")
+        return bos
+
+    def encode(self, text, *args, **kwargs):
+        if isinstance(text, str):
+            return self._encode_one(text, *args, **kwargs)
+        elif isinstance(text, list):
+            return [self._encode_one(t, *args, **kwargs) for t in text]
+        else:
+            raise ValueError(f"Invalid input type: {type(text)}")
+
+    def __call__(self, *args, **kwargs):
+        return self.encode(*args, **kwargs)
+
+    def decode(self, ids):
+        return self.tokenizer.decode(ids, skip_special_tokens=False)
+
+    def save(self, tokenizer_dir):
+        # save the tokenizer to disk
+        os.makedirs(tokenizer_dir, exist_ok=True)
+        tokenizer_path = os.path.join(tokenizer_dir, "tokenizer.json")
+        self.tokenizer.save(tokenizer_path)
+        print(f"Saved tokenizer to {tokenizer_path}")
+
+# -----------------------------------------------------------------------------
+# Tokenizer based on rustbpe + tiktoken combo
+import pickle
+import rustbpe
+import tiktoken
+
+class RustBPETokenizer:
+    """Light wrapper around tiktoken (for efficient inference) but train with rustbpe"""
+
+    def __init__(self, enc, bos_token):
+        self.enc = enc
+        self.bos_token_id = self.encode_special(bos_token)
+
+    @classmethod
+    def train_from_iterator(cls, text_iterator, vocab_size):
+        # 1) train using rustbpe
+        tokenizer = rustbpe.Tokenizer()
+        # the special tokens are inserted later in __init__, we don't train them here
+        vocab_size_no_special = vocab_size - len(SPECIAL_TOKENS)
+        assert vocab_size_no_special >= 256, f"vocab_size_no_special must be at least 256, got {vocab_size_no_special}"
+        tokenizer.train_from_iterator(text_iterator, vocab_size_no_special, pattern=SPLIT_PATTERN)
+        # 2) construct the associated tiktoken encoding for inference
+        pattern = tokenizer.get_pattern()
+        mergeable_ranks_list = tokenizer.get_mergeable_ranks()
+        mergeable_ranks = {bytes(k): v for k, v in mergeable_ranks_list}
+        tokens_offset = len(mergeable_ranks)
+        special_tokens = {name: tokens_offset + i for i, name in enumerate(SPECIAL_TOKENS)}
+        enc = tiktoken.Encoding(
+            name="rustbpe",
+            pat_str=pattern,
+            mergeable_ranks=mergeable_ranks, # dict[bytes, int] (token bytes -> merge priority rank)
+            special_tokens=special_tokens, # dict[str, int] (special token name -> token id)
+        )
+        return cls(enc, "<|bos|>")
+
+    @classmethod
+    def from_directory(cls, tokenizer_dir):
+        pickle_path = os.path.join(tokenizer_dir, "tokenizer.pkl")
+        with open(pickle_path, "rb") as f:
+            enc = pickle.load(f)
+        return cls(enc, "<|bos|>")
+
+    @classmethod
+    def from_pretrained(cls, tiktoken_name):
+        # https://github.com/openai/tiktoken/blob/eedc8563/tiktoken_ext/openai_public.py
+        enc = tiktoken.get_encoding(tiktoken_name)
+        # tiktoken calls the special document delimiter token "<|endoftext|>"
+        # yes this is confusing because this token is almost always PREPENDED to the beginning of the document
+        # it most often is used to signal the start of a new sequence to the LLM during inference etc.
+        # so in nanoChat we always use "<|bos|>" short for "beginning of sequence", but historically it is often called "<|endoftext|>".
+        return cls(enc, "<|endoftext|>")
+
+    def get_vocab_size(self):
+        return self.enc.n_vocab
+
+    def get_special_tokens(self):
+        return self.enc.special_tokens_set
+
+    def id_to_token(self, id):
+        return self.enc.decode([id])
+
+    @lru_cache(maxsize=32)
+    def encode_special(self, text):
+        return self.enc.encode_single_token(text)
+
+    def get_bos_token_id(self):
+        return self.bos_token_id
+
+    def encode(self, text, prepend=None, append=None, num_threads=8):
+        # text can be either a string or a list of strings
+
+        if prepend is not None:
+            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
+        if append is not None:
+            append_id = append if isinstance(append, int) else self.encode_special(append)
+
+        if isinstance(text, str):
+            ids = self.enc.encode_ordinary(text)
+            if prepend is not None:
+                ids.insert(0, prepend_id) # TODO: slightly inefficient here? :( hmm
+            if append is not None:
+                ids.append(append_id)
+        elif isinstance(text, list):
+            ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)
+            if prepend is not None:
+                for ids_row in ids:
+                    ids_row.insert(0, prepend_id) # TODO: same
+            if append is not None:
+                for ids_row in ids:
+                    ids_row.append(append_id)
+        else:
+            raise ValueError(f"Invalid input type: {type(text)}")
+
+        return ids
+
+    def __call__(self, *args, **kwargs):
+        return self.encode(*args, **kwargs)
+
+    def decode(self, ids):
+        return self.enc.decode(ids)
+
+    def save(self, tokenizer_dir):
+        # save the encoding object to disk
+        os.makedirs(tokenizer_dir, exist_ok=True)
+        pickle_path = os.path.join(tokenizer_dir, "tokenizer.pkl")
+        with open(pickle_path, "wb") as f:
+            pickle.dump(self.enc, f)
+        print(f"Saved tokenizer encoding to {pickle_path}")
+
+    def render_conversation(self, conversation, max_tokens=2048):
+        """
+        Tokenize a single Chat conversation (which we call a "doc" or "document" here).
+        Returns:
+        - ids: list[int] is a list of token ids of this rendered conversation
+        - mask: list[int] of same length, mask = 1 for tokens that the Assistant is expected to train on.
+        """
+        # ids, masks that we will return and a helper function to help build them up.
+        ids, mask = [], []
+        def add_tokens(token_ids, mask_val):
+            if isinstance(token_ids, int):
+                token_ids = [token_ids]
+            ids.extend(token_ids)
+            mask.extend([mask_val] * len(token_ids))
+
+        # sometimes the first message is a system message...
+        # => just merge it with the second (user) message
+        if conversation["messages"][0]["role"] == "system":
+            # some conversation surgery is necessary here for now...
+            conversation = copy.deepcopy(conversation) # avoid mutating the original
+            messages = conversation["messages"]
+            assert messages[1]["role"] == "user", "System message must be followed by a user message"
+            messages[1]["content"] = messages[0]["content"] + "\n\n" + messages[1]["content"]
+            messages = messages[1:]
+        else:
+            messages = conversation["messages"]
+        assert len(messages) >= 1, f"Conversation has less than 1 message: {messages}"
+
+        # fetch all the special tokens we need
+        bos = self.get_bos_token_id()
+        user_start, user_end = self.encode_special("<|user_start|>"), self.encode_special("<|user_end|>")
+        assistant_start, assistant_end = self.encode_special("<|assistant_start|>"), self.encode_special("<|assistant_end|>")
+        python_start, python_end = self.encode_special("<|python_start|>"), self.encode_special("<|python_end|>")
+        output_start, output_end = self.encode_special("<|output_start|>"), self.encode_special("<|output_end|>")
+
+        # now we can tokenize the conversation
+        add_tokens(bos, 0)
+        for i, message in enumerate(messages):
+
+            # some sanity checking here around assumptions, to prevent footguns
+            must_be_from = "user" if i % 2 == 0 else "assistant"
+            assert message["role"] == must_be_from, f"Message {i} is from {message['role']} but should be from {must_be_from}"
+
+            # content can be either a simple string or a list of parts (e.g. containing tool calls)
+            content = message["content"]
+
+            if message["role"] == "user":
+                assert isinstance(content, str), "User messages are simply expected to be strings"
+                value_ids = self.encode(content)
+                add_tokens(user_start, 0)
+                add_tokens(value_ids, 0)
+                add_tokens(user_end, 0)
+            elif message["role"] == "assistant":
+                add_tokens(assistant_start, 0)
+                if isinstance(content, str):
+                    # simple string => simply add the tokens
+                    value_ids = self.encode(content)
+                    add_tokens(value_ids, 1)
+                elif isinstance(content, list):
+                    for part in content:
+                        value_ids = self.encode(part["text"])
+                        if part["type"] == "text":
+                            # string part => simply add the tokens
+                            add_tokens(value_ids, 1)
+                        elif part["type"] == "python":
+                            # python tool call => add the tokens inside <|python_start|> and <|python_end|>
+                            add_tokens(python_start, 1)
+                            add_tokens(value_ids, 1)
+                            add_tokens(python_end, 1)
+                        elif part["type"] == "python_output":
+                            # python output => add the tokens inside <|output_start|> and <|output_end|>
+                            # none of these tokens are supervised because the tokens come from Python at test time
+                            add_tokens(output_start, 0)
+                            add_tokens(value_ids, 0)
+                            add_tokens(output_end, 0)
+                        else:
+                            raise ValueError(f"Unknown part type: {part['type']}")
+                else:
+                    raise ValueError(f"Unknown content type: {type(content)}")
+                add_tokens(assistant_end, 1)
+
+        # truncate to max_tokens tokens MAX (helps prevent OOMs)
+        ids = ids[:max_tokens]
+        mask = mask[:max_tokens]
+        return ids, mask
+
+    def visualize_tokenization(self, ids, mask, with_token_id=False):
+        """Small helper function useful in debugging: visualize the tokenization of render_conversation"""
+        RED = '\033[91m'
+        GREEN = '\033[92m'
+        RESET = '\033[0m'
+        GRAY = '\033[90m'
+        tokens = []
+        for i, (token_id, mask_val) in enumerate(zip(ids, mask)):
+            token_str = self.decode([token_id])
+            color = GREEN if mask_val == 1 else RED
+            tokens.append(f"{color}{token_str}{RESET}")
+            if with_token_id:
+                tokens.append(f"{GRAY}({token_id}){RESET}")
+        return '|'.join(tokens)
+
+    def render_for_completion(self, conversation):
+        """
+        Used during Reinforcement Learning. In that setting, we want to
+        render the conversation priming the Assistant for a completion.
+        Unlike the Chat SFT case, we don't need to return the mask.
+        """
+        # We have some surgery to do: we need to pop the last message (of the Assistant)
+        conversation = copy.deepcopy(conversation) # avoid mutating the original
+        messages = conversation["messages"]
+        assert messages[-1]["role"] == "assistant", "Last message must be from the Assistant"
+        messages.pop() # remove the last message (of the Assistant) inplace
+
+        # Now tokenize the conversation
+        ids, mask = self.render_conversation(conversation)
+
+        # Finally, to prime the Assistant for a completion, append the Assistant start token
+        assistant_start = self.encode_special("<|assistant_start|>")
+        ids.append(assistant_start)
+        return ids
+
+# -----------------------------------------------------------------------------
+# nanochat-specific convenience functions
+
+def get_tokenizer():
+    from nanochat.common import get_base_dir
+    base_dir = get_base_dir()
+    tokenizer_dir = os.path.join(base_dir, "tokenizer")
+    # return HuggingFaceTokenizer.from_directory(tokenizer_dir)
+    return RustBPETokenizer.from_directory(tokenizer_dir)
+
+def get_token_bytes(device="cpu"):
+    import torch
+    from nanochat.common import get_base_dir
+    base_dir = get_base_dir()
+    tokenizer_dir = os.path.join(base_dir, "tokenizer")
+    token_bytes_path = os.path.join(tokenizer_dir, "token_bytes.pt")
+    assert os.path.exists(token_bytes_path), f"Token bytes not found at {token_bytes_path}? It gets written by tok_train.py"
+    with open(token_bytes_path, "rb") as f:
+        token_bytes = torch.load(f, map_location=device)
+    return token_bytes

nanochat/ui.html

@@ -0,0 +1,562 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0, viewport-fit=cover">
+    <title>NanoChat</title>
+    <link rel="icon" type="image/svg+xml" href="/logo.svg">
+    <style>
+        :root {
+            color-scheme: light;
+        }
+
+        * {
+            box-sizing: border-box;
+        }
+
+        body {
+            font-family: ui-sans-serif, -apple-system, system-ui, "Segoe UI", Helvetica, "Apple Color Emoji", Arial, sans-serif, "Segoe UI Emoji", "Segoe UI Symbol";
+            background-color: #ffffff;
+            color: #111827;
+            min-height: 100dvh;
+            margin: 0;
+            display: flex;
+            flex-direction: column;
+        }
+
+        .header {
+            background-color: #ffffff;
+            padding: 1.25rem 1.5rem;
+        }
+
+        .header-left {
+            display: flex;
+            align-items: center;
+            gap: 0.75rem;
+        }
+
+        .header-logo {
+            height: 32px;
+            width: auto;
+        }
+
+        .header h1 {
+            font-size: 1.25rem;
+            font-weight: 600;
+            margin: 0;
+            color: #111827;
+        }
+
+        .new-conversation-btn {
+            width: 32px;
+            height: 32px;
+            padding: 0;
+            border: 1px solid #e5e7eb;
+            border-radius: 0.5rem;
+            background-color: #ffffff;
+            color: #6b7280;
+            cursor: pointer;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            transition: all 0.2s ease;
+        }
+
+        .new-conversation-btn:hover {
+            background-color: #f3f4f6;
+            border-color: #d1d5db;
+            color: #374151;
+        }
+
+        .chat-container {
+            flex: 1;
+            overflow-y: auto;
+            background-color: #ffffff;
+        }
+
+        .chat-wrapper {
+            max-width: 48rem;
+            margin: 0 auto;
+            padding: 2rem 1.5rem 3rem;
+            display: flex;
+            flex-direction: column;
+            gap: 0.75rem;
+        }
+
+        .message {
+            display: flex;
+            justify-content: flex-start;
+            margin-bottom: 0.5rem;
+            color: #0d0d0d;
+        }
+
+        .message.assistant {
+            justify-content: flex-start;
+        }
+
+        .message.user {
+            justify-content: flex-end;
+        }
+
+        .message-content {
+            white-space: pre-wrap;
+            line-height: 1.6;
+            max-width: 100%;
+        }
+
+        .message.assistant .message-content {
+            background: transparent;
+            border: none;
+            padding: 0.25rem 0;
+            cursor: pointer;
+            border-radius: 0.5rem;
+            padding: 0.5rem;
+            margin-left: -0.5rem;
+            transition: background-color 0.2s ease;
+        }
+
+        .message.assistant .message-content:hover {
+            background-color: #f9fafb;
+        }
+
+        .message.user .message-content {
+            background-color: #f3f4f6;
+            border-radius: 1.25rem;
+            padding: 0.8rem 1rem;
+            max-width: 65%;
+            cursor: pointer;
+            transition: background-color 0.2s ease;
+        }
+
+        .message.user .message-content:hover {
+            background-color: #e5e7eb;
+        }
+
+        .message.console .message-content {
+            font-family: 'Monaco', 'Menlo', 'Ubuntu Mono', 'Consolas', 'Courier New', monospace;
+            font-size: 0.875rem;
+            background-color: #fafafa;
+            padding: 0.75rem 1rem;
+            color: #374151;
+            max-width: 80%;
+        }
+
+        .input-container {
+            background-color: #ffffff;
+            padding: 1rem;
+            padding-bottom: calc(1rem + env(safe-area-inset-bottom))
+        }
+
+        .input-wrapper {
+            max-width: 48rem;
+            margin: 0 auto;
+            display: flex;
+            gap: 0.75rem;
+            align-items: flex-end;
+        }
+
+        .chat-input {
+            flex: 1;
+            padding: 0.8rem 1rem;
+            border: 1px solid #d1d5db;
+            border-radius: 0.75rem;
+            background-color: #ffffff;
+            color: #111827;
+            font-size: 1rem;
+            line-height: 1.5;
+            resize: none;
+            outline: none;
+            min-height: 54px;
+            max-height: 200px;
+            transition: border-color 0.2s ease, box-shadow 0.2s ease;
+        }
+
+        .chat-input::placeholder {
+            color: #9ca3af;
+        }
+
+        .chat-input:focus {
+            border-color: #2563eb;
+            box-shadow: 0 0 0 3px rgba(37, 99, 235, 0.1);
+        }
+
+        .send-button {
+            flex-shrink: 0;
+            padding: 0;
+            width: 54px;
+            height: 54px;
+            border: 1px solid #111827;
+            border-radius: 0.75rem;
+            background-color: #111827;
+            color: #ffffff;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            cursor: pointer;
+            transition: background-color 0.2s ease, border-color 0.2s ease, color 0.2s ease;
+        }
+
+        .send-button:hover:not(:disabled) {
+            background-color: #2563eb;
+            border-color: #2563eb;
+        }
+
+        .send-button:disabled {
+            cursor: not-allowed;
+            border-color: #d1d5db;
+            background-color: #e5e7eb;
+            color: #9ca3af;
+        }
+
+        .typing-indicator {
+            display: inline-block;
+            color: #6b7280;
+            letter-spacing: 0.15em;
+        }
+
+        .typing-indicator::after {
+            content: '···';
+            animation: typing 1.4s infinite;
+        }
+
+        @keyframes typing {
+            0%, 60%, 100% { opacity: 0.2; }
+            30% { opacity: 1; }
+        }
+
+        .error-message {
+            background-color: #fee2e2;
+            border: 1px solid #fecaca;
+            color: #b91c1c;
+            padding: 0.75rem 1rem;
+            border-radius: 0.75rem;
+            margin-top: 0.5rem;
+        }
+    </style>
+</head>
+<body>
+    <div class="header">
+        <div class="header-left">
+            <button class="new-conversation-btn" onclick="newConversation()" title="New Conversation (Ctrl+Shift+N)">
+                <svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+                    <path d="M12 5v14"></path>
+                    <path d="M5 12h14"></path>
+                </svg>
+            </button>
+            <h1>nanochat</h1>
+        </div>
+    </div>
+
+    <div class="chat-container" id="chatContainer">
+        <div class="chat-wrapper" id="chatWrapper">
+            <!-- Messages will be added here -->
+        </div>
+    </div>
+
+    <div class="input-container">
+        <div class="input-wrapper">
+            <textarea
+                id="chatInput"
+                class="chat-input"
+                placeholder="Ask anything"
+                rows="1"
+                onkeydown="handleKeyDown(event)"
+            ></textarea>
+            <button id="sendButton" class="send-button" onclick="sendMessage()" disabled>
+                <svg width="22" height="22" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
+                    <path d="M22 2L11 13"></path>
+                    <path d="M22 2l-7 20-4-9-9-4 20-7z"></path>
+                </svg>
+            </button>
+        </div>
+    </div>
+
+    <script>
+        const API_URL = '';
+        const chatContainer = document.getElementById('chatContainer');
+        const chatWrapper = document.getElementById('chatWrapper');
+        const chatInput = document.getElementById('chatInput');
+        const sendButton = document.getElementById('sendButton');
+
+        let messages = [];
+        let isGenerating = false;
+        let currentTemperature = 0.8;
+        let currentTopK = 50;
+
+        chatInput.addEventListener('input', function() {
+            this.style.height = 'auto';
+            this.style.height = Math.min(this.scrollHeight, 200) + 'px';
+            sendButton.disabled = !this.value.trim() || isGenerating;
+        });
+
+        function handleKeyDown(event) {
+            if (event.key === 'Enter' && !event.shiftKey) {
+                event.preventDefault();
+                sendMessage();
+            }
+        }
+
+        document.addEventListener('keydown', function(event) {
+            // Ctrl+Shift+N for new conversation
+            if (event.ctrlKey && event.shiftKey && event.key === 'N') {
+                event.preventDefault();
+                if (!isGenerating) {
+                    newConversation();
+                }
+            }
+        });
+
+        function newConversation() {
+            messages = [];
+            chatWrapper.innerHTML = '';
+            chatInput.value = '';
+            chatInput.style.height = 'auto';
+            sendButton.disabled = false;
+            isGenerating = false;
+            chatInput.focus();
+        }
+
+        function addMessage(role, content, messageIndex = null) {
+            const messageDiv = document.createElement('div');
+            messageDiv.className = `message ${role}`;
+
+            const contentDiv = document.createElement('div');
+            contentDiv.className = 'message-content';
+            contentDiv.textContent = content;
+
+            // Add click handler for user messages to enable editing
+            if (role === 'user' && messageIndex !== null) {
+                contentDiv.setAttribute('data-message-index', messageIndex);
+                contentDiv.setAttribute('title', 'Click to edit and restart from here');
+                contentDiv.addEventListener('click', function() {
+                    if (!isGenerating) {
+                        editMessage(messageIndex);
+                    }
+                });
+            }
+
+            // Add click handler for assistant messages to enable regeneration
+            if (role === 'assistant' && messageIndex !== null) {
+                contentDiv.setAttribute('data-message-index', messageIndex);
+                contentDiv.setAttribute('title', 'Click to regenerate this response');
+                contentDiv.addEventListener('click', function() {
+                    if (!isGenerating) {
+                        regenerateMessage(messageIndex);
+                    }
+                });
+            }
+
+            messageDiv.appendChild(contentDiv);
+            chatWrapper.appendChild(messageDiv);
+
+            chatContainer.scrollTop = chatContainer.scrollHeight;
+            return contentDiv;
+        }
+
+        function editMessage(messageIndex) {
+            // Find the message in the messages array
+            if (messageIndex < 0 || messageIndex >= messages.length) return;
+
+            const messageToEdit = messages[messageIndex];
+            if (messageToEdit.role !== 'user') return;
+
+            // Copy message content to input
+            chatInput.value = messageToEdit.content;
+            chatInput.style.height = 'auto';
+            chatInput.style.height = Math.min(chatInput.scrollHeight, 200) + 'px';
+
+            // Remove this message and all subsequent messages from the array
+            messages = messages.slice(0, messageIndex);
+
+            // Remove message elements from DOM starting from messageIndex
+            const allMessages = chatWrapper.querySelectorAll('.message');
+            for (let i = messageIndex; i < allMessages.length; i++) {
+                allMessages[i].remove();
+            }
+
+            // Enable send button and focus input
+            sendButton.disabled = false;
+            chatInput.focus();
+        }
+
+        async function generateAssistantResponse() {
+            isGenerating = true;
+            sendButton.disabled = true;
+
+            const assistantContent = addMessage('assistant', '');
+            assistantContent.innerHTML = '<span class="typing-indicator"></span>';
+
+            try {
+                const response = await fetch(`${API_URL}/chat/completions`, {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json',
+                    },
+                    body: JSON.stringify({
+                        messages: messages,
+                        temperature: currentTemperature,
+                        top_k: currentTopK,
+                        max_tokens: 512
+                    }),
+                });
+
+                if (!response.ok) {
+                    throw new Error(`HTTP error! status: ${response.status}`);
+                }
+
+                const reader = response.body.getReader();
+                const decoder = new TextDecoder();
+                let fullResponse = '';
+                assistantContent.textContent = '';
+
+                while (true) {
+                    const { done, value } = await reader.read();
+                    if (done) break;
+
+                    const chunk = decoder.decode(value);
+                    const lines = chunk.split('\n');
+
+                    for (const line of lines) {
+                        if (line.startsWith('data: ')) {
+                            try {
+                                const data = JSON.parse(line.slice(6));
+                                if (data.token) {
+                                    fullResponse += data.token;
+                                    assistantContent.textContent = fullResponse;
+                                    chatContainer.scrollTop = chatContainer.scrollHeight;
+                                }
+                            } catch (e) {
+                            }
+                        }
+                    }
+                }
+
+                const assistantMessageIndex = messages.length;
+                messages.push({ role: 'assistant', content: fullResponse });
+
+                // Add click handler to regenerate this assistant message
+                assistantContent.setAttribute('data-message-index', assistantMessageIndex);
+                assistantContent.setAttribute('title', 'Click to regenerate this response');
+                assistantContent.addEventListener('click', function() {
+                    if (!isGenerating) {
+                        regenerateMessage(assistantMessageIndex);
+                    }
+                });
+
+            } catch (error) {
+                console.error('Error:', error);
+                assistantContent.innerHTML = `<div class="error-message">Error: ${error.message}</div>`;
+            } finally {
+                isGenerating = false;
+                sendButton.disabled = !chatInput.value.trim();
+            }
+        }
+
+        async function regenerateMessage(messageIndex) {
+            // Find the message in the messages array
+            if (messageIndex < 0 || messageIndex >= messages.length) return;
+
+            const messageToRegenerate = messages[messageIndex];
+            if (messageToRegenerate.role !== 'assistant') return;
+
+            // Remove this message and all subsequent messages from the array
+            messages = messages.slice(0, messageIndex);
+
+            // Remove message elements from DOM starting from messageIndex
+            const allMessages = chatWrapper.querySelectorAll('.message');
+            for (let i = messageIndex; i < allMessages.length; i++) {
+                allMessages[i].remove();
+            }
+
+            // Regenerate the assistant response
+            await generateAssistantResponse();
+        }
+
+        function handleSlashCommand(command) {
+            const parts = command.trim().split(/\s+/);
+            const cmd = parts[0].toLowerCase();
+            const arg = parts[1];
+
+            if (cmd === '/temperature') {
+                if (arg === undefined) {
+                    addMessage('console', `Current temperature: ${currentTemperature}`);
+                } else {
+                    const temp = parseFloat(arg);
+                    if (isNaN(temp) || temp < 0 || temp > 2) {
+                        addMessage('console', 'Invalid temperature. Must be between 0.0 and 2.0');
+                    } else {
+                        currentTemperature = temp;
+                        addMessage('console', `Temperature set to ${currentTemperature}`);
+                    }
+                }
+                return true;
+            } else if (cmd === '/topk') {
+                if (arg === undefined) {
+                    addMessage('console', `Current top-k: ${currentTopK}`);
+                } else {
+                    const topk = parseInt(arg);
+                    if (isNaN(topk) || topk < 1 || topk > 200) {
+                        addMessage('console', 'Invalid top-k. Must be between 1 and 200');
+                    } else {
+                        currentTopK = topk;
+                        addMessage('console', `Top-k set to ${currentTopK}`);
+                    }
+                }
+                return true;
+            } else if (cmd === '/clear') {
+                newConversation();
+                return true;
+            } else if (cmd === '/help') {
+                addMessage('console',
+                    'Available commands:\n' +
+                    '/temperature - Show current temperature\n' +
+                    '/temperature <value> - Set temperature (0.0-2.0)\n' +
+                    '/topk - Show current top-k\n' +
+                    '/topk <value> - Set top-k (1-200)\n' +
+                    '/clear - Clear conversation\n' +
+                    '/help - Show this help message'
+                );
+                return true;
+            }
+            return false;
+        }
+
+        async function sendMessage() {
+            const message = chatInput.value.trim();
+            if (!message || isGenerating) return;
+
+            // Handle slash commands
+            if (message.startsWith('/')) {
+                chatInput.value = '';
+                chatInput.style.height = 'auto';
+                handleSlashCommand(message);
+                return;
+            }
+
+            chatInput.value = '';
+            chatInput.style.height = 'auto';
+
+            const userMessageIndex = messages.length;
+            messages.push({ role: 'user', content: message });
+            addMessage('user', message, userMessageIndex);
+
+            await generateAssistantResponse();
+        }
+
+        sendButton.disabled = false;
+
+        // Autofocus the chat input on page load
+        chatInput.focus();
+
+        fetch(`${API_URL}/health`)
+            .then(response => response.json())
+            .then(data => {
+                console.log('Engine status:', data);
+            })
+            .catch(error => {
+                console.error('Engine not available:', error);
+                chatWrapper.innerHTML = '<div class="error-message">Engine not running. Please start engine.py first.</div>';
+            });
+    </script>
+</body>
+</html>

pyproject.toml

@@ -0,0 +1,77 @@
+[project]
+name = "nanochat"
+version = "0.1.0"
+description = "the minimal full-stack ChatGPT clone"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "datasets>=4.0.0",
+    "fastapi>=0.117.1",
+    "files-to-prompt>=0.6",
+    "psutil>=7.1.0",
+    "regex>=2025.9.1",
+    "setuptools>=80.9.0",
+    "tiktoken>=0.11.0",
+    "tokenizers>=0.22.0",
+    "torch>=2.8.0",
+    "uvicorn>=0.36.0",
+    "wandb>=0.21.3",
+]
+
+[build-system]
+requires = ["maturin>=1.7,<2.0"]
+build-backend = "maturin"
+
+[tool.maturin]
+module-name = "rustbpe"
+bindings = "pyo3"
+python-source = "."
+manifest-path = "rustbpe/Cargo.toml"
+
+[dependency-groups]
+dev = [
+    "maturin>=1.9.4",
+    "pytest>=8.0.0",
+]
+
+[tool.pytest.ini_options]
+markers = [
+    "slow: marks tests as slow (deselect with '-m \"not slow\"')",
+]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+
+# target torch to cuda 12.8 or CPU
+[tool.uv.sources]
+torch = [  
+    { index = "pytorch-cpu", extra = "cpu" },  
+    { index = "pytorch-cu128", extra = "gpu" },  
+]
+
+[[tool.uv.index]]  
+name = "pytorch-cpu"  
+url = "https://download.pytorch.org/whl/cpu"  
+explicit = true  
+  
+[[tool.uv.index]]  
+name = "pytorch-cu128"  
+url = "https://download.pytorch.org/whl/cu128"  
+explicit = true
+
+[project.optional-dependencies]  
+cpu = [  
+    "torch>=2.8.0",  
+]  
+gpu = [  
+    "torch>=2.8.0",  
+]  
+  
+[tool.uv]  
+conflicts = [  
+    [  
+        { extra = "cpu" },  
+        { extra = "gpu" },  
+    ],  
+]  

run1000.sh

@@ -0,0 +1,94 @@
+#!/bin/bash
+
+# The $1000 tier of nanochat
+# Designed to run end-to-end for $1000/24 ~= 41.6 hours on an 8XH100 node
+# A bit sparser on comments, see speedrun.sh for more detail
+
+# all the setup stuff
+export OMP_NUM_THREADS=1
+export NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
+mkdir -p $NANOCHAT_BASE_DIR
+command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
+[ -d ".venv" ] || uv venv
+uv sync --extra gpu
+source .venv/bin/activate
+if [ -z "$WANDB_RUN" ]; then
+    WANDB_RUN=dummy
+fi
+python -m nanochat.report reset
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+source "$HOME/.cargo/env"
+uv run maturin develop --release --manifest-path rustbpe/Cargo.toml
+curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl
+
+# train tokenizer on ~4B characters and kick off download of the rest for pretraining
+python -m nanochat.dataset -n 16
+# start downloading the rest of the shards for a total of 800 (see below why 800)
+python -m nanochat.dataset -n 800 &
+# todo: download the rest of it
+python -m scripts.tok_train --max_chars=4000000000
+python -m scripts.tok_eval
+
+# Documenting my process for determining the hyperparameters for this run1000.sh script:
+# We want a budget of approx. $1000 ~= 41.6 hours of 8XH100 compute
+# 1) I guessed the model size for this to be about depth=32
+# 2) Determine the device_batch_size that fits:
+# Running the base_train.py script with --depth=32, I saw that --device_batch_size=16
+# runs out of memory, but --device_batch_size=8 fits. Inspecting `nvidia-smi` during training,
+# I saw all GPUs were at about 78/80GB VRAM, so it just barely fits and we have good MFU at ~50%.
+# So the training script was running ok and showed:
+# Vocab size: 65,536
+# num_layers: 32
+# model_dim: 2048
+# num_heads: 16
+# num_kv_heads: 16
+# Tokens / micro-batch / rank: 8 x 2048 = 16,384
+# Tokens / micro-batch: 131,072
+# Total batch size 524,288 => gradient accumulation steps: 4
+# Number of parameters: 1,879,048,192
+# Estimated FLOPs per token: 1.207960e+10
+# Calculated number of iterations from target data:param ratio: 71,680
+# Total number of training tokens: 37,580,963,840
+# Tokens : Params ratio: 20.00
+# Total training FLOPs estimate: 4.539628e+20
+# step 00004/71680 (0.01%) | loss: 8.813754 | lrm: 1.00 | dt: 1571.88ms | tok/sec: 83,385 | mfu: 50.92 | total time: 0.00m
+# step 00005/71680 (0.01%) | loss: 8.488074 | lrm: 1.00 | dt: 1572.76ms | tok/sec: 83,338 | mfu: 50.89 | total time: 0.00m
+# ...
+# 3) validate that the runtime fits our budget:
+# The training script uses the Chinchilla scaling law to compute-optimally set #tokens = 20 * #params. In particular:
+# The script shows that we will be training for 71,680 steps, and each step takes 1.574s so:
+# estimated time to train: 71,680 * 1.574s / 60 / 60 = 31.3 hours.
+# This is OK, fits our budget, and leaves ~10 hours for midtraining and SFT and evals and maybe RL.
+# It's possible that we might even fit depth=33 or depth=34, but for now let's go along with this.
+# 4) The last thing to pay attention to is the amount of training data required for the run.
+# The script above calculated that "Total number of training tokens: 37,580,963,840"
+# The tok_eval.py script reports about ~4.8 chars/token on average for the default tokenizer settings.
+# So ~38B tokens # ~4.8 chars/token = ~185B chars.
+# Each data shard is ~250M chars, so we need ~185B / 250M ~= 740 shards.
+# For safety, I bumped that up to 800 shards, and that's why up above I used -n 800 when pre-downloading dataset shards.
+# If we didn't have enough data, the training script would loop around and do multiple epochs over the same data,
+# which would decrease model performance. Possibly 2, 3 or so epochs is ~ok, but certainly not ideal and at 10+ epochs we'd
+# start to overfit hard.
+# 5) That's it, everything else (e.g. the learning rates) is adjusted automatically by the training script.
+
+# Number of processes/GPUs to use
+NPROC_PER_NODE=8
+
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_train -- --depth=32 --device_batch_size=8 --run=$WANDB_RUN
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_loss
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_eval
+
+# midtrain
+# NOTE: ensure that we use the same device_batch_size here as the base training script.
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.mid_train -- --device_batch_size=8 --run=$WANDB_RUN
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_eval -- -i mid
+
+# sft
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_sft -- --run=$WANDB_RUN
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_eval -- -i sft
+
+# generate final report
+python -m nanochat.report generate
+
+# talk to it
+python -m scripts.chat_web

rustbpe/Cargo.lock

@@ -0,0 +1,458 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "ahash"
+version = "0.8.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75"
+dependencies = [
+ "cfg-if",
+ "getrandom",
+ "once_cell",
+ "version_check",
+ "zerocopy",
+]
+
+[[package]]
+name = "aho-corasick"
+version = "1.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916"
+dependencies = [
+ "memchr",
+]
+
+[[package]]
+name = "arc-swap"
+version = "1.7.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69f7f8c3906b62b754cd5326047894316021dcfe5a194c8ea52bdd94934a3457"
+
+[[package]]
+name = "autocfg"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
+
+[[package]]
+name = "bit-set"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "08807e080ed7f9d5433fa9b275196cfc35414f66a0c79d864dc51a0d825231a3"
+dependencies = [
+ "bit-vec",
+]
+
+[[package]]
+name = "bit-vec"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
+
+[[package]]
+name = "castaway"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a"
+dependencies = [
+ "rustversion",
+]
+
+[[package]]
+name = "cfg-if"
+version = "1.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2fd1289c04a9ea8cb22300a459a72a385d7c73d3259e2ed7dcb2af674838cfa9"
+
+[[package]]
+name = "compact_str"
+version = "0.9.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a"
+dependencies = [
+ "castaway",
+ "cfg-if",
+ "itoa",
+ "rustversion",
+ "ryu",
+ "static_assertions",
+]
+
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
+
+[[package]]
+name = "dary_heap"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "04d2cd9c18b9f454ed67da600630b021a8a80bf33f8c95896ab33aaf1c26b728"
+
+[[package]]
+name = "either"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
+
+[[package]]
+name = "equivalent"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f"
+
+[[package]]
+name = "fancy-regex"
+version = "0.16.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf04c5ec15464ace8355a7b440a33aece288993475556d461154d7a62ad9947c"
+dependencies = [
+ "bit-set",
+ "regex-automata",
+ "regex-syntax",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "r-efi",
+ "wasi",
+]
+
+[[package]]
+name = "hashbrown"
+version = "0.15.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "indexmap"
+version = "2.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2481980430f9f78649238835720ddccc57e52df14ffce1c6f37391d61b563e9"
+dependencies = [
+ "equivalent",
+ "hashbrown",
+]
+
+[[package]]
+name = "indoc"
+version = "2.0.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd"
+
+[[package]]
+name = "itoa"
+version = "1.0.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
+
+[[package]]
+name = "libc"
+version = "0.2.175"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6a82ae493e598baaea5209805c49bbf2ea7de956d50d7da0da1164f9c6d28543"
+
+[[package]]
+name = "log"
+version = "0.4.28"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "34080505efa8e45a4b816c349525ebe327ceaa8559756f0356cba97ef3bf7432"
+
+[[package]]
+name = "memchr"
+version = "2.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0"
+
+[[package]]
+name = "memoffset"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d"
+
+[[package]]
+name = "portable-atomic"
+version = "1.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f84267b20a16ea918e43c6a88433c2d54fa145c92a811b5b047ccbe153674483"
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.101"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
+dependencies = [
+ "cfg-if",
+ "indoc",
+ "libc",
+ "memoffset",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+ "unindent",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
+dependencies = [
+ "once_cell",
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-log"
+version = "0.12.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "45192e5e4a4d2505587e27806c7b710c231c40c56f3bfc19535d0bb25df52264"
+dependencies = [
+ "arc-swap",
+ "log",
+ "pyo3",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.23.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "pyo3-build-config",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "r-efi"
+version = "5.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
+
+[[package]]
+name = "rayon"
+version = "1.11.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "regex-automata"
+version = "0.4.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6b9458fa0bfeeac22b5ca447c63aaf45f28439a709ccd244698632f9aa6394d6"
+dependencies = [
+ "aho-corasick",
+ "memchr",
+ "regex-syntax",
+]
+
+[[package]]
+name = "regex-syntax"
+version = "0.8.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001"
+
+[[package]]
+name = "rustbpe"
+version = "0.1.0"
+dependencies = [
+ "ahash",
+ "compact_str",
+ "dary_heap",
+ "fancy-regex",
+ "indexmap",
+ "log",
+ "pyo3",
+ "pyo3-log",
+ "rayon",
+]
+
+[[package]]
+name = "rustversion"
+version = "1.0.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
+
+[[package]]
+name = "ryu"
+version = "1.0.20"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f"
+
+[[package]]
+name = "static_assertions"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f"
+
+[[package]]
+name = "syn"
+version = "2.0.106"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.12.16"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.18"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
+
+[[package]]
+name = "unindent"
+version = "0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
+
+[[package]]
+name = "version_check"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
+
+[[package]]
+name = "wasi"
+version = "0.14.4+wasi-0.2.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "88a5f4a424faf49c3c2c344f166f0662341d470ea185e939657aaff130f0ec4a"
+dependencies = [
+ "wit-bindgen",
+]
+
+[[package]]
+name = "wit-bindgen"
+version = "0.45.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c573471f125075647d03df72e026074b7203790d41351cd6edc96f46bcccd36"
+
+[[package]]
+name = "zerocopy"
+version = "0.8.26"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f"
+dependencies = [
+ "zerocopy-derive",
+]
+
+[[package]]
+name = "zerocopy-derive"
+version = "0.8.26"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]

rustbpe/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "rustbpe"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+dary_heap = "0.3"
+indexmap = "2.2"
+fancy-regex = "0.16.1"
+log = "0.4.28"
+pyo3 = { version = "0.23.3", features = ["extension-module"] }
+pyo3-log = "0.12.4"
+ahash = "0.8.12"
+rayon = "1.11.0"
+compact_str = "0.9.0"

rustbpe/README.md

@@ -0,0 +1,5 @@
+# rustbpe
+
+> The missing tiktoken training code
+
+A very lightweight Rust library for training a GPT tokenizer. The issue is that the inference library [tiktoken](https://github.com/openai/tiktoken) is great, but only does inference. Separately, the huggingface [tokenizers](https://github.com/huggingface/tokenizers) library does training, but it is rather bloated and really hard to navigate because it has to support all the different historical baggage of how people dealt with tokenizers over the years. More recently, I also wrote the [minbpe](https://github.com/karpathy/minbpe) library which does both training and inference, but only in inefficient Python. Basically what I really want is a non-fancy, super simple, but still relatively efficient training code for GPT tokenizer (more efficient than minbpe, much cleaner/simpler than tokenizers), and then export the trained vocab for inference with tiktoken. Does that make sense? So here we are. There are more opportunities for optimization here, I just stopped a bit early because unlike minbpe before it, rustbpe is now simple and fast enough, and not a significant bottleneck for nanochat.

rustbpe/src/lib.rs

@@ -0,0 +1,475 @@
+use std::cmp::Ordering;
+use std::collections::HashMap as StdHashMap;
+
+use dary_heap::OctonaryHeap;
+use fancy_regex::Regex;
+use pyo3::prelude::*;
+
+use ahash::{AHashMap, AHashSet};
+use compact_str::CompactString;
+use rayon::prelude::*;
+
+// Default GPT-4 style regex pattern for splitting text
+const GPT4_PATTERN: &str = r"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+";
+
+type Pair = (u32, u32);
+
+/// A Byte Pair Encoding tokenizer that matches the GPT-4 style implementation
+#[pyclass]
+pub struct Tokenizer {
+    /// Maps pairs of token IDs to their merged token ID
+    pub merges: StdHashMap<Pair, u32>,
+    /// The regex pattern used for text splitting
+    pub pattern: String,
+    /// Compiled regex for efficiency
+    compiled_pattern: Regex,
+}
+
+// ------------------------ internal helpers ------------------------
+
+#[derive(Clone, Debug)]
+struct Word {
+    ids: Vec<u32>,
+}
+
+impl Word {
+    #[inline]
+    fn new(ids: Vec<u32>) -> Self {
+        Self { ids }
+    }
+
+    #[inline]
+    fn pairs<'a>(&'a self) -> impl Iterator<Item = Pair> + 'a {
+        self.ids.windows(2).map(|w| (w[0], w[1]))
+    }
+
+    /// Merge all non-overlapping occurrences of pair -> new_id.
+    /// Returns a small Vec of local pair-count deltas for THIS word only:
+    ///   -1 for removed pairs, +1 for newly created pairs.
+    ///
+    /// NOTE: this version deliberately avoids a HashMap in the hot loop.
+    fn merge_pair(&mut self, pair: Pair, new_id: u32) -> Vec<(Pair, i32)> {
+        let (a, b) = pair;
+        let n = self.ids.len();
+        if n < 2 {
+            return Vec::new();
+        }
+
+        let mut out: Vec<u32> = Vec::with_capacity(n);
+        let mut deltas: Vec<(Pair, i32)> = Vec::with_capacity(6);
+
+        let mut i = 0;
+        while i < n {
+            if i + 1 < n && self.ids[i] == a && self.ids[i + 1] == b {
+                let left = out.last().copied();
+                let right = if i + 2 < n { Some(self.ids[i + 2]) } else { None };
+
+                // remove old pairs
+                if let Some(x) = left {
+                    deltas.push(((x, a), -1));
+                    deltas.push(((x, new_id), 1));
+                }
+                deltas.push(((a, b), -1));
+                if let Some(y) = right {
+                    deltas.push(((b, y), -1));
+                    deltas.push(((new_id, y), 1));
+                }
+
+                // write merged token
+                out.push(new_id);
+                i += 2; // skip 'a' and 'b'
+            } else {
+                out.push(self.ids[i]);
+                i += 1;
+            }
+        }
+
+        self.ids = out;
+        deltas
+    }
+}
+
+#[derive(Debug, Eq)]
+struct MergeJob {
+    pair: Pair,
+    count: u64,
+    /// set of word indices where this pair may occur and needs processing
+    pos: AHashSet<usize>,
+}
+
+impl PartialEq for MergeJob {
+    fn eq(&self, other: &Self) -> bool {
+        self.count == other.count && self.pair == other.pair
+    }
+}
+
+impl PartialOrd for MergeJob {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for MergeJob {
+    fn cmp(&self, other: &Self) -> Ordering {
+        // Max-heap by count; tie-break to ascending pair order (deterministic)
+        if self.count != other.count {
+            self.count.cmp(&other.count)
+        } else {
+            // ascending order on the pair when counts tie
+            other.pair.cmp(&self.pair)
+        }
+    }
+}
+
+#[inline]
+fn count_pairs_parallel(
+    words: &[Word],
+    counts: &[i32],
+) -> (AHashMap<Pair, i32>, AHashMap<Pair, AHashSet<usize>>) {
+    words
+        .par_iter()
+        .enumerate()
+        .map(|(i, w)| {
+            let mut local_pc: AHashMap<Pair, i32> = AHashMap::new();
+            let mut local_wtu: AHashMap<Pair, AHashSet<usize>> = AHashMap::new();
+            if w.ids.len() >= 2 && counts[i] != 0 {
+                for (a, b) in w.pairs() {
+                    *local_pc.entry((a, b)).or_default() += counts[i];
+                    local_wtu.entry((a, b)).or_default().insert(i);
+                }
+            }
+            (local_pc, local_wtu)
+        })
+        .reduce(
+            || (AHashMap::new(), AHashMap::new()),
+            |(mut acc_pc, mut acc_wtu), (pc, wtu)| {
+                for (k, v) in pc {
+                    *acc_pc.entry(k).or_default() += v;
+                }
+                for (k, s) in wtu {
+                    acc_wtu.entry(k).or_default().extend(s);
+                }
+                (acc_pc, acc_wtu)
+            },
+        )
+}
+
+// ------------------------ END helpers ------------------------
+
+impl Tokenizer {
+
+    /// Core incremental BPE training given unique words and their counts.
+    /// `words`: one entry per unique chunk (Vec<u32> of token-ids/bytes).
+    /// `counts`: same length as `words`, count per chunk.
+    fn train_core_incremental(&mut self, mut words: Vec<Word>, counts: Vec<i32>, vocab_size: u32) {
+        assert!(vocab_size >= 256, "vocab_size must be at least 256");
+        let num_merges = vocab_size - 256;
+        log::info!("Starting BPE training: {} merges to compute", num_merges);
+        self.merges.clear();
+
+        // ---- Initial pair_counts and where_to_update (parallel) ----
+        log::info!("Computing initial pair counts from {} unique sequences", words.len());
+        let (mut pair_counts, mut where_to_update) = count_pairs_parallel(&words, &counts);
+
+        // ---- Build heap ----
+        log::info!("Building heap with {} unique pairs", pair_counts.len());
+        let mut heap = OctonaryHeap::with_capacity(pair_counts.len());
+        for (pair, pos) in where_to_update.drain() {
+            let c = *pair_counts.get(&pair).unwrap_or(&0);
+            if c > 0 {
+                heap.push(MergeJob {
+                    pair,
+                    count: c as u64,
+                    pos,
+                });
+            }
+        }
+
+        // ---- Merge loop ----
+        log::info!("Starting merge loop");
+        let mut merges_done = 0u32;
+        let mut last_log_percent = 0u32;
+
+        while merges_done < num_merges {
+            let Some(mut top) = heap.pop() else { break; };
+
+            // Lazy refresh
+            let current = *pair_counts.get(&top.pair).unwrap_or(&0);
+            if top.count != current as u64 {
+                top.count = current as u64;
+                if top.count > 0 {
+                    heap.push(top);
+                }
+                continue;
+            }
+            if top.count == 0 {
+                break;
+            }
+
+            // Record merge
+            let new_id = 256 + merges_done;
+            self.merges.insert(top.pair, new_id);
+
+            // Merge this pair in all words where it occurs
+            let mut local_pos_updates: AHashMap<Pair, AHashSet<usize>> = AHashMap::new();
+            for &word_idx in &top.pos {
+                // Apply merge to this word and collect pair-count deltas
+                let changes = words[word_idx].merge_pair(top.pair, new_id);
+                // Update global pair counts based on this word's count
+                for (pair, delta) in changes {
+                    let delta_total = delta * counts[word_idx];
+                    if delta_total != 0 {
+                        *pair_counts.entry(pair).or_default() += delta_total;
+                        if delta > 0 {
+                            local_pos_updates.entry(pair).or_default().insert(word_idx);
+                        }
+                    }
+                }
+            }
+
+            // Add the updated pair counts back to the heap
+            for (pair, pos) in local_pos_updates {
+                let cnt = *pair_counts.get(&pair).unwrap_or(&0);
+                if cnt > 0 {
+                    heap.push(MergeJob {
+                        pair,
+                        count: cnt as u64,
+                        pos,
+                    });
+                }
+            }
+
+            merges_done += 1;
+
+            // Log progress every 1%
+            let current_percent = (merges_done * 100) / num_merges;
+            if current_percent > last_log_percent {
+                log::info!(
+                    "Progress: {}% ({}/{} merges) - Last merge: {:?} -> {} (frequency: {})",
+                    current_percent, merges_done, num_merges, top.pair, new_id, top.count
+                );
+                last_log_percent = current_percent;
+            }
+        }
+
+        log::info!("Finished training: {} merges completed", merges_done);
+    }
+}
+
+/// Public methods for the Tokenizer class that will be exposed to Python.
+#[pymethods]
+impl Tokenizer {
+    /// Create a new Tokenizer
+    #[new]
+    pub fn new() -> Self {
+        Self {
+            merges: StdHashMap::new(),
+            pattern: String::new(),
+            compiled_pattern: Regex::new("").expect("Empty regex should be valid"),
+        }
+    }
+
+    /// Train from a streaming iterator (parallel ingestion).
+    /// We refill a Rust Vec<String> buffer under the GIL, then release the GIL
+    /// to do the heavy splitting and counting **in parallel** with rayon.
+    #[pyo3(signature = (iterator, vocab_size, buffer_size=8192, pattern=None))]
+    #[pyo3(text_signature = "(self, iterator, vocab_size, buffer_size=8192, pattern=None)")]
+    pub fn train_from_iterator(
+        &mut self,
+        py: pyo3::Python<'_>,
+        iterator: &pyo3::Bound<'_, pyo3::PyAny>,
+        vocab_size: u32,
+        buffer_size: usize,
+        pattern: Option<String>,
+    ) -> PyResult<()> {
+        // Use provided pattern or default to GPT-4 pattern
+        let pattern_str = pattern.unwrap_or_else(|| GPT4_PATTERN.to_string());
+
+        // Update the stored pattern and compile it
+        self.pattern = pattern_str.clone();
+        self.compiled_pattern = Regex::new(&pattern_str)
+            .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("Invalid regex pattern: {}", e)))?;
+
+        // Prepare a true Python iterator object
+        let py_iter: pyo3::Py<pyo3::PyAny> = unsafe {
+            pyo3::Py::from_owned_ptr_or_err(py, pyo3::ffi::PyObject_GetIter(iterator.as_ptr()))?
+        };
+
+        // Global chunk counts
+        let mut counts: AHashMap<CompactString, i32> = AHashMap::new();
+
+        // Temporary buffer we refill under the GIL
+        let mut buf: Vec<String> = Vec::with_capacity(buffer_size);
+
+        log::info!("Processing sequences from iterator (buffer_size: {})", buffer_size);
+        let mut total_sequences = 0u64;
+
+        // Helper: refill `buf` with up to `buffer_size` strings from the Python iterator.
+        // Returns Ok(true) if the iterator is exhausted, Ok(false) otherwise.
+        let refill = |buf: &mut Vec<String>| -> PyResult<bool> {
+            pyo3::Python::with_gil(|py| {
+                buf.clear();
+                let it = py_iter.bind(py);
+                loop {
+                    if buf.len() >= buffer_size {
+                        return Ok(false);
+                    }
+                    // next(it)
+                    let next_obj = unsafe {
+                        pyo3::Bound::from_owned_ptr_or_opt(py, pyo3::ffi::PyIter_Next(it.as_ptr()))
+                    };
+                    match next_obj {
+                        Some(obj) => {
+                            let s: String = obj.extract()?;
+                            buf.push(s);
+                        }
+                        None => {
+                            if pyo3::PyErr::occurred(py) {
+                                return Err(pyo3::PyErr::fetch(py));
+                            } else {
+                                return Ok(true); // exhausted
+                            }
+                        }
+                    }
+                }
+            })
+        };
+
+        // Stream ingestion loop: refill under GIL, process without GIL (parallel)
+        loop {
+            let exhausted = refill(&mut buf)?;
+            if buf.is_empty() && exhausted {
+                break;
+            }
+
+            total_sequences += buf.len() as u64;
+
+            let pattern = self.compiled_pattern.clone();
+            let local: AHashMap<CompactString, i32> = py.allow_threads(|| {
+                buf.par_iter()
+                    .map(|s| {
+                        let mut m: AHashMap<CompactString, i32> = AHashMap::new();
+                        for mat in pattern.find_iter(s) {
+                            let piece = mat.expect("regex match failed").as_str();
+                            *m.entry(CompactString::from(piece)).or_default() += 1;
+                        }
+                        m
+                    })
+                    .reduce(
+                        || AHashMap::new(),
+                        |mut a, b| {
+                            for (k, v) in b {
+                                *a.entry(k).or_default() += v;
+                            }
+                            a
+                        },
+                    )
+            });
+
+            // Merge local into global (single-threaded)
+            for (k, v) in local {
+                *counts.entry(k).or_default() += v;
+            }
+
+            if exhausted {
+                break;
+            }
+        }
+        log::info!("Processed {} sequences total, {} unique", total_sequences, counts.len());
+
+        // Materialize words & counts
+        let mut words = Vec::with_capacity(counts.len());
+        let mut cvec = Vec::with_capacity(counts.len());
+        for (chunk, c) in counts.into_iter() {
+            words.push(Word::new(chunk.as_bytes().iter().map(|&b| b as u32).collect()));
+            cvec.push(c);
+        }
+
+        self.train_core_incremental(words, cvec, vocab_size);
+        Ok(())
+    }
+
+    /// Return the regex pattern
+    pub fn get_pattern(&self) -> String {
+        self.pattern.clone()
+    }
+
+    /// Return the mergeable ranks (token bytes -> token id / rank)
+    pub fn get_mergeable_ranks(&self) -> Vec<(Vec<u8>, u32)> {
+        let mut mergeable_ranks = Vec::new();
+
+        // Build vocabulary incrementally from low to high token IDs
+        let mut token_bytes: Vec<Vec<u8>> = (0..256_u32).map(|i| vec![i as u8]).collect();
+
+        for (i, bytes) in token_bytes.iter().enumerate() {
+            mergeable_ranks.push((bytes.clone(), i as u32));
+        }
+
+        // Sort merges by token id (so we can reconstruct bytes progressively)
+        let mut sorted_merges: Vec<_> = self.merges.iter().collect();
+        sorted_merges.sort_by_key(|&(_, &token_id)| token_id);
+
+        for (&pair, &merged_id) in sorted_merges {
+            let (left, right) = pair;
+            let mut merged_bytes = token_bytes[left as usize].clone();
+            merged_bytes.extend(&token_bytes[right as usize]);
+
+            if token_bytes.len() <= merged_id as usize {
+                token_bytes.resize(merged_id as usize + 1, Vec::new());
+            }
+            token_bytes[merged_id as usize] = merged_bytes.clone();
+
+            mergeable_ranks.push((merged_bytes, merged_id));
+        }
+
+        mergeable_ranks
+    }
+
+    /// Encode a string into token IDs
+    pub fn encode(&self, text: &str) -> Vec<u32> {
+        let mut all_ids = Vec::new();
+
+        // Split text using the regex pattern
+        for m in self.compiled_pattern.find_iter(text) {
+            let chunk = m.expect("regex match failed").as_str();
+
+            // Convert chunk to bytes then to u32 IDs
+            let mut ids: Vec<u32> = chunk.bytes().map(|b| b as u32).collect();
+
+            // Apply merges iteratively
+            while ids.len() >= 2 {
+                // Find the best pair to merge
+                let mut best_pair: Option<(usize, Pair, u32)> = None;
+
+                for i in 0..ids.len() - 1 {
+                    let pair: Pair = (ids[i], ids[i + 1]);
+                    if let Some(&new_id) = self.merges.get(&pair) {
+                        if best_pair.is_none() || new_id < best_pair.unwrap().2 {
+                            best_pair = Some((i, pair, new_id));
+                        }
+                    }
+                }
+
+                // If we found a pair to merge, apply it
+                if let Some((idx, _pair, new_id)) = best_pair {
+                    ids[idx] = new_id;
+                    ids.remove(idx + 1);
+                } else {
+                    // No more merges possible
+                    break;
+                }
+            }
+
+            all_ids.extend(ids);
+        }
+
+        all_ids
+    }
+}
+
+#[pymodule]
+fn rustbpe(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    pyo3_log::init(); // forwards Rust `log` to Python's `logging`
+    m.add_class::<Tokenizer>()?;
+    Ok(())
+}

scripts/base_eval.py

@@ -0,0 +1,212 @@
+"""
+Evaluate the CORE metric for a given model.
+
+Run on a single GPU:
+python -m scripts.base_eval
+
+Run with torchrun on e.g. 8 GPUs:
+torchrun --nproc_per_node=8 -m scripts.base_eval
+
+The script will print the CORE metric to the console.
+"""
+import os
+import csv
+import time
+import json
+import yaml
+import shutil
+import random
+import zipfile
+import tempfile
+from contextlib import nullcontext
+
+import torch
+
+from nanochat.common import compute_init, compute_cleanup, print0, get_base_dir, autodetect_device_type, download_file_with_lock
+from nanochat.tokenizer import HuggingFaceTokenizer
+from nanochat.checkpoint_manager import load_model
+from nanochat.core_eval import evaluate_task
+
+# -----------------------------------------------------------------------------
+# nanochat specific function dealing with I/O etc.
+
+# ~162MB of data needed to evaluate the CORE metric
+EVAL_BUNDLE_URL = "https://karpathy-public.s3.us-west-2.amazonaws.com/eval_bundle.zip"
+
+def place_eval_bundle(file_path):
+    # here file_path is the path to the eval_bundle.zip file
+    # we need to unzip it and place it in the base directory
+    base_dir = get_base_dir()
+    eval_bundle_dir = os.path.join(base_dir, "eval_bundle")
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with zipfile.ZipFile(file_path, 'r') as zip_ref:
+            zip_ref.extractall(tmpdir)
+        extracted_bundle_dir = os.path.join(tmpdir, "eval_bundle")
+        shutil.move(extracted_bundle_dir, eval_bundle_dir)
+    print0(f"Placed eval_bundle directory at {eval_bundle_dir}")
+
+def evaluate_model(model, tokenizer, device, max_per_task=-1):
+    """
+    Evaluate a base model on the CORE benchmark.
+    - max_per_task: crop the data to this many examples per task for testing (-1 = disable)
+    """
+    # Load config and task metadata
+    base_dir = get_base_dir()
+    eval_bundle_dir = os.path.join(base_dir, "eval_bundle")
+    # Download the eval bundle to disk (and unzip if needed)
+    if not os.path.exists(eval_bundle_dir):
+        download_file_with_lock(EVAL_BUNDLE_URL, "eval_bundle.zip", postprocess_fn=place_eval_bundle)
+    config_path = os.path.join(eval_bundle_dir, "core.yaml")
+    data_base_path = os.path.join(eval_bundle_dir, "eval_data")
+    eval_meta_data = os.path.join(eval_bundle_dir, "eval_meta_data.csv")
+    with open(config_path, 'r', encoding='utf-8') as f:
+        config = yaml.safe_load(f)
+    tasks = config['icl_tasks']
+
+    # Load random baseline values from eval metadata
+    random_baselines = {}
+    with open(eval_meta_data, 'r', encoding='utf-8') as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            task_name = row['Eval Task']
+            random_baseline = row['Random baseline']
+            random_baselines[task_name] = float(random_baseline)
+
+    # Evaluate each task
+    results = {}
+    centered_results = {}
+    for task in tasks:
+        start_time = time.time()
+        label = task['label']
+        task_meta = {
+            'task_type': task['icl_task_type'],
+            'dataset_uri': task['dataset_uri'],
+            'num_fewshot': task['num_fewshot'][0],
+            'continuation_delimiter': task.get('continuation_delimiter', ' ')
+        }
+        print0(f"Evaluating: {label} ({task_meta['num_fewshot']}-shot, type: {task_meta['task_type']})... ", end='')
+
+        # Load data for this task
+        data_path = os.path.join(data_base_path, task_meta['dataset_uri'])
+        with open(data_path, 'r', encoding='utf-8') as f:
+            data = [json.loads(line.strip()) for line in f]
+
+        # shuffle the data because in many cases it appears ordered but we want
+        # the ability to only run a subset of the data for debugging purposes etc.
+        shuffle_rng = random.Random(1337)
+        shuffle_rng.shuffle(data)
+        if max_per_task > 0:
+            data = data[:max_per_task]
+
+        # run the evaluation for this task
+        accuracy = evaluate_task(model, tokenizer, data, device, task_meta)
+
+        results[label] = accuracy
+        random_baseline = random_baselines[label]
+        centered_result = (accuracy - 0.01 * random_baseline) / (1.0 - 0.01 * random_baseline)
+        centered_results[label] = centered_result
+        end_time = time.time()
+        print0(f"accuracy: {accuracy:.4f} | centered: {centered_result:.4f} | time: {end_time - start_time:.2f}s")
+
+    core_metric = sum(centered_results.values()) / len(centered_results)
+    out = {
+        "results": results,
+        "centered_results": centered_results,
+        "core_metric": core_metric
+    }
+    return out
+
+# -----------------------------------------------------------------------------
+# HuggingFace loading utilities and light wrappers for a model
+
+class ModelWrapper:
+    """Lightweight wrapper for a HuggingFace model"""
+    def __init__(self, model, max_seq_len=None):
+        self.model = model
+        self.max_seq_len = max_seq_len
+
+    def __call__(self, input_ids):
+        outputs = self.model(input_ids)
+        logits = outputs.logits
+        return logits
+
+def load_hf_model(hf_path: str, device):
+    print0(f"Loading model from: {hf_path}")
+    # Load the model
+    from transformers import AutoModelForCausalLM
+    model = AutoModelForCausalLM.from_pretrained(hf_path)
+    model.to(device)
+    model.eval()
+    max_seq_len = 1024 if "openai-community/gpt2" in hf_path else None
+    model = ModelWrapper(model, max_seq_len=max_seq_len)
+    # Load the tokenizer
+    tokenizer = HuggingFaceTokenizer.from_pretrained(hf_path)
+    return model, tokenizer
+
+# -----------------------------------------------------------------------------
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--hf-path', type=str, default=None, help='HuggingFace model path to evaluate')
+    parser.add_argument('--max-per-task', type=int, default=-1, help='Max examples per task to evaluate (-1 = disable)')
+    args = parser.parse_args()
+
+    # distributed / precision setup
+    device_type = autodetect_device_type()
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+    autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
+
+    # Load model and tokenizer from command line or from file system
+    if args.hf_path is not None:
+        # atm assume that if a path is given, it's a huggingface model path
+        hf_path = args.hf_path
+        print0(f"Loading huggingface model from: {hf_path}")
+        model, tokenizer = load_hf_model(hf_path, device)
+        model_name = hf_path # just for logging
+        model_slug = hf_path.replace("/", "-") # for the output csv file
+    else:
+        # load a local model from the file system
+        model, tokenizer, meta = load_model("base", device, phase="eval")
+        model_name = f"base_model (step {meta['step']})" # just for logging
+        model_slug = f"base_model_{meta['step']:06d}" # for the output csv file
+
+    # Evaluate the model
+    with autocast_ctx:
+        out = evaluate_model(model, tokenizer, device, max_per_task=args.max_per_task)
+
+    # Write out the results to a csv file
+    core_metric = None
+    centered_results = {}
+    if ddp_rank == 0:
+        base_dir = get_base_dir()
+        output_csv_path = os.path.join(base_dir, "base_eval", f"{model_slug}.csv")
+        os.makedirs(os.path.dirname(output_csv_path), exist_ok=True)
+        results = out["results"]
+        centered_results = out["centered_results"]
+        core_metric = out["core_metric"]
+        with open(output_csv_path, 'w', encoding='utf-8', newline='') as f:
+            f.write(f"{'Task':<35}, {'Accuracy':<10}, {'Centered':<10}\n")
+            for label in results:
+                f.write(f"{label:<35}, {results[label]:<10.6f}, {centered_results[label]:<10.6f}\n")
+            f.write(f"{'CORE':<35}, {'':<10}, {core_metric:<10.6f}\n")
+        # Print the content of the csv file to console too
+        print0("="*80)
+        print0(f"Model: {model_name}")
+        print0("="*80)
+        with open(output_csv_path, 'r', encoding='utf-8') as f:
+            print0(f.read())
+
+    # Log to report
+    from nanochat.report import get_report
+    get_report().log(section="Base model evaluation", data=[
+        {
+            "Model": model_name,
+            "CORE metric": core_metric,
+        },
+        centered_results, # the full table
+    ])
+
+    compute_cleanup()
+
+if __name__ == "__main__":
+    main()

scripts/base_loss.py

@@ -0,0 +1,79 @@
+"""
+Loads a checkpoint, and:
+- Evaluates the loss on a larger chunk of train/val splits
+- Samples from the model
+
+Example run as:
+torchrun --standalone --nproc_per_node=8 -m scripts.base_loss
+"""
+import os
+from contextlib import nullcontext
+import torch
+from nanochat.checkpoint_manager import load_model
+from nanochat.common import compute_init, print0, compute_cleanup, autodetect_device_type
+from nanochat.dataloader import tokenizing_distributed_data_loader
+from nanochat.tokenizer import get_token_bytes
+from nanochat.loss_eval import evaluate_bpb
+from nanochat.engine import Engine
+
+# Configuration
+device_batch_size = 32
+split_tokens = 20*524288  # number of tokens to evaluate per split
+model_tag = None # optional model tag for the output directory name
+model_step = None # optional model step for the output directory name
+device_type = "" # cuda|cpu|mps (empty => autodetect)
+exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file
+
+# Load the base model and the tokenizer
+device_type = autodetect_device_type() if device_type == "" else device_type
+ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+model, tokenizer, meta = load_model("base", device, phase="eval", model_tag=model_tag, step=model_step)
+sequence_len = meta["model_config"]["sequence_len"] # could be arbitrary really
+autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
+
+# Evaluate the loss on each split
+tokens_per_step = device_batch_size * sequence_len * ddp_world_size
+assert split_tokens % tokens_per_step == 0, "split_tokens must be divisible by tokens_per_step"
+steps = split_tokens // tokens_per_step
+token_bytes = get_token_bytes(device=device)
+bpb_results = {}
+for split_name in ["train", "val"]:
+    loader = tokenizing_distributed_data_loader(device_batch_size, sequence_len, split_name, device=device)
+    with autocast_ctx:
+        bpb = evaluate_bpb(model, loader, steps, token_bytes)
+    print0(f"{split_name} bpb: {bpb:.4f}")
+    bpb_results[split_name] = bpb
+
+# Master process also samples from the model
+samples = []
+if ddp_rank == 0:
+    prompts = [
+        "The capital of France is",
+        "The chemical symbol of gold is",
+        "If yesterday was Friday, then tomorrow will be",
+        "The opposite of hot is",
+        "The planets of the solar system are:",
+        "My favorite color is",
+        "If 5*x + 3 = 13, then x is",
+    ]
+    engine = Engine(model, tokenizer)
+    for prompt in prompts:
+        tokens = tokenizer(prompt, prepend="<|bos|>")
+        with autocast_ctx:
+            sample, _ = engine.generate_batch(tokens, num_samples=1, max_tokens=16, temperature=0)
+        sample_str = tokenizer.decode(sample[0])
+        print0(sample_str)
+        samples.append(sample_str)
+
+# Log to report
+from nanochat.report import get_report
+get_report().log(section="Base model loss", data=[
+    {
+        "train bpb": bpb_results["train"],
+        "val bpb": bpb_results["val"],
+    },
+    {f"sample {i}": sample for i, sample in enumerate(samples)},
+])
+
+# Cleanup
+compute_cleanup()

scripts/base_train.py

@@ -0,0 +1,355 @@
+"""
+Train model. Run as:
+
+python base_train.py
+
+or distributed as:
+
+torchrun --nproc_per_node=8 base_train.py
+
+If you are only on CPU/Macbook, you'll want to train a much much smaller LLM. Example:
+python -m scripts.base_train --depth=4 --max_seq_len=512 --device_batch_size=1 --eval_tokens=512 --core_metric_every=-1 --total_batch_size=512 --num_iterations=20
+"""
+
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+import time
+from contextlib import nullcontext
+
+import wandb
+import torch
+
+from nanochat.gpt import GPT, GPTConfig
+from nanochat.dataloader import tokenizing_distributed_data_loader
+from nanochat.common import compute_init, compute_cleanup, print0, DummyWandb, print_banner, get_base_dir, autodetect_device_type
+from nanochat.tokenizer import get_tokenizer, get_token_bytes
+from nanochat.checkpoint_manager import save_checkpoint
+from nanochat.loss_eval import evaluate_bpb
+from nanochat.engine import Engine
+from scripts.base_eval import evaluate_model
+print_banner()
+
+# -----------------------------------------------------------------------------
+# User settings
+run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb)
+# Runtime
+device_type = "" # cuda|cpu|mps (empty => autodetect good device type default, in order: CUDA > MPS > CPU)
+# Model architecture
+depth = 20 # the depth of the Transformer model to train, rest of the kwargs are derived
+max_seq_len = 2048 # max context length
+# Training horizon. Only one of these 3 will be used, in this order of precedence.
+num_iterations = -1 # explicit number of steps of the optimization (-1 = disable)
+target_flops = -1.0 # calculate num_iterations to reach target_flops. Useful for scaling laws experiments (-1 = disable)
+target_param_data_ratio = 20 # calculate num_iterations to maintain fixed data:param ratio (Chinchilla=20) (-1 = disable)
+# Optimization
+device_batch_size = 32 # per-device batch size (set to not OOM)
+total_batch_size = 524288 # total desired batch size, in #tokens
+embedding_lr = 0.2 # learning rate for the embedding parameters (Adam)
+unembedding_lr = 0.004 # learning rate for the unembedding parameters (Adam)
+weight_decay = 0.0 # weight decay for the embedding/unembedding parameters (Adam)
+matrix_lr = 0.02 # learning rate for the matrix parameters (Muon)
+grad_clip = 1.0 # gradient clipping value (0.0 = disabled)
+warmup_ratio = 0.0 # ratio of iterations for LR warmup
+warmdown_ratio = 0.2 # ratio of iterations for LR warmdown
+final_lr_frac = 0.0 # final LR is this fraction of the initial LR
+# Evaluation
+eval_every = 250 # every how many steps to evaluate the model for val bpb
+eval_tokens = 20*524288 # number of tokens to evaluate val loss on
+core_metric_every = 2000 # every how many steps to evaluate the core metric (-1 = disable)
+core_metric_max_per_task = 500 # examples per task in estimating the core metric
+sample_every = 2000 # every how many steps to sample from the model
+# Output
+model_tag = "" # optionally override the model tag for the output checkpoint directory name
+# now allow CLI to override the settings via the configurator lol
+config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file
+user_config = {k: globals()[k] for k in config_keys} # will be useful for logging
+# -----------------------------------------------------------------------------
+
+# Compute init
+device_type = autodetect_device_type() if device_type == "" else device_type
+ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
+autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
+synchronize = torch.cuda.synchronize if device_type == "cuda" else lambda: None
+get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else lambda: 0
+
+# wandb logging init
+use_dummy_wandb = run == "dummy" or not master_process
+wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat", name=run, config=user_config)
+
+# Tokenizer will be useful for evaluation, also we need the vocab size
+tokenizer = get_tokenizer()
+token_bytes = get_token_bytes(device=device)
+vocab_size = tokenizer.get_vocab_size()
+print0(f"Vocab size: {vocab_size:,}")
+
+# Model kwargs are derived from the desired depth of the model
+num_layers = depth
+model_dim = depth * 64 # aspect ratio 64 (usually this is varied from 64 -> 128 as model size increases)
+num_heads = max(1, (model_dim + 127) // 128) # head dim 128 (the division here is ceil div)
+num_kv_heads = num_heads # default is 1:1 GQA (Group Query Attention) ratio (i.e. GQA is disabled)
+print0(f"num_layers: {num_layers}")
+print0(f"model_dim: {model_dim}")
+print0(f"num_heads: {num_heads}")
+print0(f"num_kv_heads: {num_kv_heads}")
+
+# Optimizer / data / training length related hyperparameters
+# figure out the needed gradient accumulation to reach the desired total batch size
+tokens_per_fwdbwd = device_batch_size * max_seq_len # tokens per iteration for a single rank
+world_tokens_per_fwdbwd = tokens_per_fwdbwd * ddp_world_size # total tokens per iteration for all ranks
+assert total_batch_size % world_tokens_per_fwdbwd == 0
+grad_accum_steps = total_batch_size // world_tokens_per_fwdbwd
+print0(f"Tokens / micro-batch / rank: {device_batch_size} x {max_seq_len} = {tokens_per_fwdbwd:,}")
+print0(f"Tokens / micro-batch: {world_tokens_per_fwdbwd:,}")
+print0(f"Total batch size {total_batch_size:,} => gradient accumulation steps: {grad_accum_steps}")
+# -----------------------------------------------------------------------------
+# Initialize the Model
+model_config_kwargs = dict(sequence_len=max_seq_len, vocab_size=vocab_size, n_layer=num_layers, n_head=num_heads, n_kv_head=num_kv_heads, n_embd=model_dim)
+with torch.device("meta"):
+    model_config = GPTConfig(**model_config_kwargs)
+    model = GPT(model_config)
+model.to_empty(device=device)
+model.init_weights()
+orig_model = model # original, uncompiled model, for saving raw model state_dict
+model = torch.compile(model, dynamic=False) # TODO: dynamic True/False think through
+num_params = sum(p.numel() for p in model.parameters())
+print0(f"Number of parameters: {num_params:,}")
+num_flops_per_token = model.estimate_flops()
+print0(f"Estimated FLOPs per token: {num_flops_per_token:e}")
+
+# Calculate number of iterations. Either it is given, or from target flops, or from target data:param ratio (in that order)
+assert num_iterations > 0 or target_param_data_ratio > 0 or target_flops > 0
+if num_iterations > 0:
+    print0(f"Using user-provided number of iterations: {num_iterations:,}")
+elif target_flops > 0:
+    # calculate the number of iterations from the target flops
+    num_iterations = round(target_flops / (num_flops_per_token * total_batch_size))
+    print0(f"Calculated number of iterations from target FLOPs: {num_iterations:,}")
+elif target_param_data_ratio > 0:
+    # calculate the number of iterations from the target param data ratio
+    target_tokens = target_param_data_ratio * num_params
+    num_iterations = target_tokens // total_batch_size
+    print0(f"Calculated number of iterations from target data:param ratio: {num_iterations:,}")
+else:
+    raise ValueError("No training horizon specified")
+total_tokens = total_batch_size * num_iterations
+print0(f"Total number of training tokens: {total_tokens:,}")
+print0(f"Tokens : Params ratio: {total_batch_size * num_iterations / num_params:.2f}") # Chinchilla is ~20
+print0(f"Total training FLOPs estimate: {num_flops_per_token * total_tokens:e}")
+
+# -----------------------------------------------------------------------------
+# Initialize the Optimizer (Muon for Linear layers, AdamW for embedding and lm_head)
+optimizers = model.setup_optimizers(unembedding_lr=unembedding_lr, embedding_lr=embedding_lr, matrix_lr=matrix_lr, weight_decay=weight_decay)
+adamw_optimizer, muon_optimizer = optimizers
+
+# Initialize the DataLoaders for train/val
+base_dir = get_base_dir()
+tokens_dir = os.path.join(base_dir, "tokenized_data")
+train_loader = tokenizing_distributed_data_loader(device_batch_size, max_seq_len, split="train", device=device)
+build_val_loader = lambda: tokenizing_distributed_data_loader(device_batch_size, max_seq_len, split="val", device=device)
+x, y = next(train_loader) # kick off load of the very first batch of data
+
+# -----------------------------------------------------------------------------
+# Set up hyperparameter schedulers
+
+# Learning rate scheduler
+def get_lr_multiplier(it):
+    warmup_iters = round(warmup_ratio * num_iterations)
+    warmdown_iters = round(warmdown_ratio * num_iterations)
+    if it < warmup_iters:
+        return (it + 1) / warmup_iters
+    elif it <= num_iterations - warmdown_iters:
+        return 1.0
+    else:
+        progress = (num_iterations - it) / warmdown_iters
+        return progress * 1.0 + (1 - progress) * final_lr_frac
+
+# Momentum scheduler for Muon optimizer
+def get_muon_momentum(it):
+    frac = min(it / 300, 1)
+    momentum = (1 - frac) * 0.85 + frac * 0.95
+    return momentum
+
+# -----------------------------------------------------------------------------
+# Training loop
+min_val_bpb = float("inf")
+smooth_train_loss = 0 # EMA of training loss
+ema_beta = 0.9 # EMA decay factor
+total_training_time = 0 # total wall-clock time of training
+# note that we run +1 steps only so that we can eval and save at the end
+for step in range(num_iterations + 1):
+    last_step = step == num_iterations
+    flops_so_far = num_flops_per_token * total_batch_size * step
+
+    # once in a while: evaluate the val bpb (all ranks participate)
+    if last_step or step % eval_every == 0:
+        model.eval()
+        val_loader = build_val_loader()
+        eval_steps = eval_tokens // (device_batch_size * max_seq_len * ddp_world_size)
+        with autocast_ctx:
+            val_bpb = evaluate_bpb(model, val_loader, eval_steps, token_bytes)
+        print0(f"Step {step:05d} | Validation bpb: {val_bpb:.4f}")
+        if val_bpb < min_val_bpb:
+            min_val_bpb = val_bpb
+        wandb_run.log({
+            "step": step,
+            "total_training_flops": flops_so_far,
+            "total_training_time": total_training_time,
+            "val/bpb": val_bpb,
+        })
+        model.train()
+
+    # once in a while: estimate the CORE metric (all ranks participate)
+    # use the original uncompiled model because the inputs keep changing shape
+    results = {}
+    if core_metric_every > 0 and (last_step or (step > 0 and step % core_metric_every == 0)):
+        model.eval()
+        with autocast_ctx:
+            results = evaluate_model(orig_model, tokenizer, device, max_per_task=core_metric_max_per_task)
+        print0(f"Step {step:05d} | CORE metric: {results['core_metric']:.4f}")
+        wandb_run.log({
+            "step": step,
+            "total_training_flops": flops_so_far,
+            "core_metric": results["core_metric"],
+            "centered_results": results["centered_results"],
+        })
+        model.train()
+
+    # once in a while: sample from the model (only on master process)
+    # use the original uncompiled model because the inputs keep changing shape
+    if master_process and (last_step or (step > 0 and step % sample_every == 0)):
+        model.eval()
+        prompts = [
+            "The capital of France is",
+            "The chemical symbol of gold is",
+            "If yesterday was Friday, then tomorrow will be",
+            "The opposite of hot is",
+            "The planets of the solar system are:",
+            "My favorite color is",
+            "If 5*x + 3 = 13, then x is",
+        ]
+        engine = Engine(orig_model, tokenizer) # use orig_model to avoid recompilation
+        for prompt in prompts:
+            tokens = tokenizer(prompt, prepend="<|bos|>")
+            with autocast_ctx:
+                sample, _ = engine.generate_batch(tokens, num_samples=1, max_tokens=16, temperature=0)
+            print0(tokenizer.decode(sample[0]))
+        model.train()
+
+    # save checkpoint at the end of the run (only on master process)
+    if master_process and last_step:
+        output_dirname = model_tag if model_tag else f"d{depth}" # e.g. d12
+        checkpoint_dir = os.path.join(base_dir, "base_checkpoints", output_dirname)
+        save_checkpoint(
+            checkpoint_dir,
+            step,
+            orig_model.state_dict(),
+            [opt.state_dict() for opt in optimizers], # TODO: make sure saving across ranks is done correctly
+            {
+                "step": step,
+                "val_bpb": val_bpb, # loss at last step
+                "model_config": model_config_kwargs,
+                "user_config": user_config, # inputs to the training script
+                "device_batch_size": device_batch_size,
+                "max_seq_len": max_seq_len,
+            }
+        )
+
+    if last_step:
+        break
+
+    # -------------------------------------------------------------------------
+    # single training step
+    # evaluate the gradient
+    synchronize()
+    t0 = time.time()
+    for micro_step in range(grad_accum_steps):
+        with autocast_ctx:
+            loss = model(x, y)
+        train_loss = loss.detach() # for logging
+        loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here
+        loss.backward()
+        x, y = next(train_loader) # prefetch the next batch while the GPU is busy with forward/backward
+    # gradient clipping
+    grad_clip_enabled = grad_clip > 0.0
+    if grad_clip_enabled:
+        grad_norm_tensor = torch.nn.utils.clip_grad_norm_(orig_model.parameters(), grad_clip)
+        grad_norm = grad_norm_tensor.item() # GPU tensor -> CPU float (note: cpu-gpu sync point)
+    # step the optimizers
+    lrm = get_lr_multiplier(step)
+    for opt in optimizers:
+        for group in opt.param_groups:
+            group["lr"] = group["initial_lr"] * lrm
+    muon_momentum = get_muon_momentum(step)
+    for group in muon_optimizer.param_groups:
+        group["momentum"] = muon_momentum
+    for opt in optimizers:
+        opt.step()
+    model.zero_grad(set_to_none=True)
+    synchronize()
+    t1 = time.time()
+    dt = t1 - t0
+    # -------------------------------------------------------------------------
+
+    # logging
+    smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss.item() # EMA the training loss
+    debiased_smooth_loss = smooth_train_loss / (1 - ema_beta**(step + 1)) # debias the EMA
+    pct_done = 100 * step / num_iterations
+    tok_per_sec = int(total_batch_size / dt)
+    flops_per_sec = num_flops_per_token * total_batch_size / dt
+    promised_flops_per_sec_h100 = 989e12 * ddp_world_size # bfloat16 H100 SXM and without 2:4 sparsity
+    mfu = 100 * flops_per_sec / promised_flops_per_sec_h100 # in %
+    if step > 10:
+        total_training_time += dt # only count the time after the first 10 steps
+    print_grad_norm = f" grad norm: {grad_norm:.4f} |" if grad_clip_enabled else ""
+    print0(f"step {step:05d}/{num_iterations:05d} ({pct_done:.2f}%) | loss: {debiased_smooth_loss:.6f} |{print_grad_norm} lrm: {lrm:.2f} | dt: {dt * 1000:.2f}ms | tok/sec: {tok_per_sec:,} | mfu: {mfu:.2f} | total time: {total_training_time/60:.2f}m")
+    if step % 100 == 0:
+        log_data = {
+            "step": step,
+            "total_training_flops": flops_so_far,
+            "total_training_time": total_training_time,
+            "train/loss": debiased_smooth_loss,
+            "train/lrm": lrm,
+            "train/dt": dt,
+            "train/tok_per_sec": tok_per_sec,
+            "train/mfu": mfu,
+        }
+        if grad_clip_enabled:
+            log_data["train/grad_norm"] = grad_norm
+        wandb_run.log(log_data)
+
+# print a few more stats
+print0(f"Peak memory usage: {get_max_memory() / 1024 / 1024:.2f}MiB")
+print0(f"Total training time: {total_training_time/60:.2f}m")
+print0(f"Minimum validation bpb: {min_val_bpb:.4f}")
+
+# Log to report
+from nanochat.report import get_report
+get_report().log(section="Base model training", data=[
+    user_config, # CLI args
+    { # stats about the training setup
+        "Number of parameters": num_params,
+        "Number of FLOPs per token": f"{num_flops_per_token:e}",
+        "Calculated number of iterations": num_iterations,
+        "Number of training tokens": total_tokens,
+        "Tokens : Params ratio": total_batch_size * num_iterations / num_params,
+        "DDP world size": ddp_world_size,
+        "warmup_ratio": warmup_ratio,
+        "warmdown_ratio": warmdown_ratio,
+        "final_lr_frac": final_lr_frac,
+    },
+    { # stats about training outcomes
+        "Minimum validation bpb": min_val_bpb,
+        "Final validation bpb": val_bpb,
+        "CORE metric estimate": results.get("core_metric", None),
+        "MFU %": f"{mfu:.2f}%",
+        "Total training flops": f"{flops_so_far:e}",
+        "Total training time": f"{total_training_time/60:.2f}m",
+        "Peak memory usage": f"{get_max_memory() / 1024 / 1024:.2f}MiB",
+    }
+])
+
+# cleanup
+wandb_run.finish() # wandb run finish
+compute_cleanup()

scripts/chat_cli.py

@@ -0,0 +1,105 @@
+"""
+New and upgraded chat mode because a lot of the code has changed since the last one.
+
+Intended to be run single GPU only atm:
+python -m scripts.chat_cli -i mid
+"""
+import argparse
+import torch
+from nanochat.common import compute_init, autodetect_device_type
+from contextlib import nullcontext
+from nanochat.engine import Engine
+from nanochat.checkpoint_manager import load_model
+
+parser = argparse.ArgumentParser(description='Chat with the model')
+parser.add_argument('-i', '--source', type=str, default="sft", help="Source of the model: sft|mid|rl")
+parser.add_argument('-g', '--model-tag', type=str, default=None, help='Model tag to load')
+parser.add_argument('-s', '--step', type=int, default=None, help='Step to load')
+parser.add_argument('-p', '--prompt', type=str, default='', help='Prompt the model, get a single response back')
+parser.add_argument('-t', '--temperature', type=float, default=0.6, help='Temperature for generation')
+parser.add_argument('-k', '--top-k', type=int, default=50, help='Top-k sampling parameter')
+parser.add_argument('--device-type', type=str, default='', choices=['cuda', 'cpu', 'mps'], help='Device type for evaluation: cuda|cpu|mps. empty => autodetect')
+parser.add_argument('-d', '--dtype', type=str, default='bfloat16', choices=['float32', 'bfloat16'])
+args = parser.parse_args()
+
+# Init the model and tokenizer
+
+device_type = autodetect_device_type() if args.device_type == "" else args.device_type
+ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+ptdtype = torch.float32 if args.dtype == 'float32' else torch.bfloat16
+autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext()
+model, tokenizer, meta = load_model(args.source, device, phase="eval", model_tag=args.model_tag, step=args.step)
+
+# Special tokens for the chat state machine
+bos = tokenizer.get_bos_token_id()
+user_start, user_end = tokenizer.encode_special("<|user_start|>"), tokenizer.encode_special("<|user_end|>")
+assistant_start, assistant_end = tokenizer.encode_special("<|assistant_start|>"), tokenizer.encode_special("<|assistant_end|>")
+
+# Create Engine for efficient generation
+engine = Engine(model, tokenizer)
+
+print("\nNanoChat Interactive Mode")
+print("-" * 50)
+print("Type 'quit' or 'exit' to end the conversation")
+print("Type 'clear' to start a new conversation")
+print("-" * 50)
+
+conversation_tokens = [bos]
+
+while True:
+
+    if args.prompt:
+        # Get the prompt from the launch command
+        user_input = args.prompt
+    else:
+        # Get the prompt interactively from the console
+        try:
+            user_input = input("\nUser: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nGoodbye!")
+            break
+
+    # Handle special commands
+    if user_input.lower() in ['quit', 'exit']:
+        print("Goodbye!")
+        break
+
+    if user_input.lower() == 'clear':
+        conversation_tokens = [bos]
+        print("Conversation cleared.")
+        continue
+
+    if not user_input:
+        continue
+
+    # Add User message to the conversation
+    conversation_tokens.append(user_start)
+    conversation_tokens.extend(tokenizer.encode(user_input))
+    conversation_tokens.append(user_end)
+
+    # Kick off the assistant
+    conversation_tokens.append(assistant_start)
+    generate_kwargs = {
+        "num_samples": 1,
+        "max_tokens": 256,
+        "temperature": args.temperature,
+        "top_k": args.top_k,
+    }
+    response_tokens = []
+    print("\nAssistant: ", end="", flush=True)
+    with autocast_ctx:
+        for token_column, token_masks in engine.generate(conversation_tokens, **generate_kwargs):
+            token = token_column[0] # pop the batch dimension (num_samples=1)
+            response_tokens.append(token)
+            token_text = tokenizer.decode([token])
+            print(token_text, end="", flush=True)
+    print()
+    # we have to ensure that the assistant end token is the last token
+    # so even if generation ends due to max tokens, we have to append it to the end
+    if response_tokens[-1] != assistant_end:
+        response_tokens.append(assistant_end)
+    conversation_tokens.extend(response_tokens)
+
+    # In the prompt mode, we only want a single response and exit
+    if args.prompt:
+        break

scripts/chat_eval.py

@@ -0,0 +1,257 @@
+"""
+Evaluate the Chat model.
+All the generic code lives here, and all the evlauation-specific
+code lives in nanochat directory and is imported from here.
+
+Example runs:
+python -m scripts.chat_eval -a ARC-Easy
+torchrun --nproc_per_node=8 -m scripts.chat_eval -- -a ARC-Easy
+"""
+
+import argparse
+from functools import partial
+from contextlib import nullcontext
+
+import torch
+import torch.distributed as dist
+
+from nanochat.common import compute_init, compute_cleanup, get_dist_info, print0, autodetect_device_type
+from nanochat.checkpoint_manager import load_model
+from nanochat.engine import Engine
+
+from tasks.humaneval import HumanEval
+from tasks.mmlu import MMLU
+from tasks.arc import ARC
+from tasks.gsm8k import GSM8K
+from tasks.spellingbee import SpellingBee
+
+# -----------------------------------------------------------------------------
+# Generative evaluation loop (we go one problem at a time, sample, evaluate)
+
+def run_generative_eval(task_object, tokenizer, model, engine, num_samples, max_new_tokens, temperature, top_k, max_problems=None):
+
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
+    device = model.get_device()
+
+    num_problems = len(task_object) if max_problems is None else min(len(task_object), max_problems)
+
+    # Run the evaluation
+    num_passed, total = 0, 0
+    for i in range(ddp_rank, num_problems, ddp_world_size):
+        conversation = task_object[i]
+
+        # Tokenize the prompt
+        encoded_prompt = tokenizer.render_for_completion(conversation)
+        # Get the completions
+        results, _ = engine.generate_batch(
+            encoded_prompt,
+            num_samples=num_samples,
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+        )
+        # Decode the completions as text
+        prefix_length = len(encoded_prompt)
+        completions = [tokenizer.decode(result_tokens[prefix_length:]) for result_tokens in results]
+        # Evaluate success criteria
+        outcomes = [task_object.evaluate(conversation, completion) for completion in completions]
+        passed = any(outcomes)
+
+        # Keep stats
+        total += 1
+        num_passed += int(passed)
+
+        # Logging (overwrite the same line in the console)
+        print(f"\r\033[KRank {ddp_rank} | {num_passed}/{total} ({100*num_passed/total:.2f}%)", end='', flush=True)
+
+    # Finish the in-place progress line with a newline before final summary
+    print()
+
+    # Aggregate results across all ranks
+    if ddp:
+        num_passed_tensor = torch.tensor([num_passed], dtype=torch.long, device=device)
+        total_tensor = torch.tensor([total], dtype=torch.long, device=device)
+        dist.all_reduce(num_passed_tensor, op=dist.ReduceOp.SUM)
+        dist.all_reduce(total_tensor, op=dist.ReduceOp.SUM)
+        num_passed = num_passed_tensor.item()
+        total = total_tensor.item()
+
+    print0("=" * 50)
+    print0(f"Final: {num_passed}/{total} ({100*num_passed/total:.2f}%)")
+
+    # Return the accuracy
+    return num_passed/total
+
+# -----------------------------------------------------------------------------
+# Categorical evaluation loop
+# A lot easier because we don't have to sample. Therefore, we can actually go
+# batches at a time and just check the logits for correct answer choices.
+
+def run_categorical_eval(task_object, tokenizer, model, batch_size, max_problems=None):
+
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info()
+    device = model.get_device()
+    bos = tokenizer.get_bos_token_id() # use BOS as pad token is ok, these positions are ignored
+
+    # We'll process batches of independent problems at a time because there is no sampling needed
+    num_problems = len(task_object) if max_problems is None else min(len(task_object), max_problems)
+    ceil_div = lambda x, y: -(-x // y)
+    num_batches = ceil_div(num_problems, batch_size)
+
+    # Run the evaluation
+    letter_to_id_cache = {} # many letters will repeat often, let's save the tokenizer some work
+    num_passed, total = 0, 0
+    for i in range(ddp_rank, num_batches, ddp_world_size):
+        i0, i1 = i * batch_size, min((i + 1) * batch_size, num_problems)
+
+        # Prepare the batch of problems. They might all be of different length, so we pad/collate them.
+        conversations = [task_object[ii] for ii in range(i0, i1)]
+        prompt_ids = [tokenizer.render_for_completion(conversation) for conversation in conversations] # TODO: remake the way this works
+        max_length = max(len(ids) for ids in prompt_ids)
+        answer_time_positions = [len(ids) - 1 for ids in prompt_ids] # where the last token is (and the predicted answer)
+        padded_prompt_ids = [ids + [bos] * (max_length - len(ids)) for ids in prompt_ids]
+        prompt_ids = torch.tensor(padded_prompt_ids, dtype=torch.long, device=device)
+
+        # Get the logits for the whole batch of conversations in parallel (efficiency win here)
+        with torch.no_grad():
+            logits = model(prompt_ids) # (B, T, V)
+
+        # Focus on the available answer on just the letters corresponding to choices
+        # Note that this helps the evaluation a lot because it specifically narrows the focus to only the available letters
+        # The much harder alternative would be to just generate from the Assistant and check if it responded with the correct
+        # letter (e.g. A, B, C, D), but evaluations typically make the task easier in this way.
+        for idx, conversation in enumerate(conversations):
+            # get the token ids of all the available letters of this problem
+            letters = conversation['letters']
+            letter_ids = []
+            for letter in letters:
+                if not letter in letter_to_id_cache:
+                    encoded_letter = tokenizer.encode(letter)
+                    assert len(encoded_letter) == 1, "Each letter must be a single token"
+                    letter_to_id_cache[letter] = encoded_letter[0]
+                letter_ids.append(letter_to_id_cache[letter])
+            # focus logits just down to the answer position and the available letters of the answer
+            answer_pos = answer_time_positions[idx]
+            focus_logits = logits[idx, answer_pos, letter_ids]
+            # get the argmax letter (the predicted answer)
+            argmax_letter_id = focus_logits.argmax(dim=-1).item()
+            predicted_letter = letters[argmax_letter_id]
+            # evaluate the outcome
+            outcome = task_object.evaluate(conversation, predicted_letter)
+            num_passed += int(outcome)
+            total += 1
+
+    # Aggregate results across all ranks
+    if ddp:
+        num_passed_tensor = torch.tensor([num_passed], dtype=torch.long, device=device)
+        total_tensor = torch.tensor([total], dtype=torch.long, device=device)
+        dist.all_reduce(num_passed_tensor, op=dist.ReduceOp.SUM)
+        dist.all_reduce(total_tensor, op=dist.ReduceOp.SUM)
+        num_passed = num_passed_tensor.item()
+        total = total_tensor.item()
+
+    average = num_passed/total
+    print0(f"Final: {num_passed}/{total} ({100*average:.2f}%)")
+    return average
+
+# -----------------------------------------------------------------------------
+
+def run_chat_eval(task_name, model, tokenizer, engine,
+                   batch_size=1, num_samples=1, max_new_tokens=512, temperature=0.0, top_k=50,
+                   max_problems=None):
+    # Create the evaluation object
+    task_module = {
+        'HumanEval': HumanEval,
+        'MMLU': partial(MMLU, subset="all", split="test"),
+        'ARC-Easy': partial(ARC, subset="ARC-Easy", split="test"),
+        'ARC-Challenge': partial(ARC, subset="ARC-Challenge", split="test"),
+        'GSM8K': partial(GSM8K, subset="main", split="test"),
+        'SpellingBee': partial(SpellingBee, size=256, split="test"),
+    }[task_name]
+    task_object = task_module()
+    # Run the evaluation
+    if task_object.eval_type == 'generative':
+        acc = run_generative_eval(task_object, tokenizer, model, engine, num_samples, max_new_tokens, temperature, top_k, max_problems=max_problems)
+    elif task_object.eval_type == 'categorical':
+        acc = run_categorical_eval(task_object, tokenizer, model, batch_size, max_problems=max_problems)
+    else:
+        raise ValueError(f"Unsupported task evaluation type: {task_object.eval_type}")
+    return acc
+
+# -----------------------------------------------------------------------------
+if __name__ == "__main__":
+
+    # Parse command-line arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--source', type=str, required=True, help="Source of the model: sft|mid|rl")
+    parser.add_argument('-a', '--task-name', type=str, default=None, help="Task name. Default = all tasks. Use | to split multiple tasks.")
+    parser.add_argument('-d', '--dtype', type=str, default='bfloat16', choices=['float32', 'bfloat16'])
+    parser.add_argument('-t', '--temperature', type=float, default=0.0)
+    parser.add_argument('-m', '--max-new-tokens', type=int, default=512)
+    parser.add_argument('-n', '--num-samples', type=int, default=1)
+    parser.add_argument('-k', '--top-k', type=int, default=50)
+    parser.add_argument('-b', '--batch-size', type=int, default=8, help='Batch size for categorical evaluation')
+    parser.add_argument('-g', '--model-tag', type=str, default=None, help='Model tag to load')
+    parser.add_argument('-s', '--step', type=int, default=None, help='Step to load')
+    parser.add_argument('-x', '--max-problems', type=int, default=None, help='Max problems to evaluate')
+    parser.add_argument('--device-type', type=str, default='', choices=['cuda', 'cpu', 'mps'], help='Device type for evaluation: cuda|cpu|mps. empty => autodetect')
+    args = parser.parse_args()
+
+    device_type = autodetect_device_type() if args.device_type == "" else args.device_type
+    ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+    ptdtype = torch.float32 if args.dtype == 'float32' else torch.bfloat16
+    autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext()
+
+    model, tokenizer, meta = load_model(args.source, device, phase="eval", model_tag=args.model_tag, step=args.step)
+    engine = Engine(model, tokenizer)
+
+    # Get the tasks to evaluate on
+    all_tasks = ['ARC-Easy', 'ARC-Challenge', 'MMLU', 'GSM8K', 'HumanEval', 'SpellingBee']
+    baseline_accuracies = {
+        'ARC-Easy': 0.25, # multiple choice 1 of 4 => 25%
+        'ARC-Challenge': 0.25, # multiple choice 1 of 4 => 25%
+        'MMLU': 0.25, # multiple choice 1 of 4 => 25%
+        'GSM8K': 0.0, # open-ended => 0%
+        'HumanEval': 0.0, # open-ended => 0%
+        'SpellingBee': 0.0, # open-ended => 0%
+    }
+    task_names = all_tasks if args.task_name is None else args.task_name.split('|')
+
+    # Run all the task evaluations sequentially
+    results = {}
+    for task_name in task_names:
+        with autocast_ctx:
+            acc = run_chat_eval(
+                task_name,
+                model, tokenizer, engine,
+                batch_size=args.batch_size,
+                num_samples=args.num_samples,
+                max_new_tokens=args.max_new_tokens,
+                temperature=args.temperature,
+                top_k=args.top_k,
+                max_problems=args.max_problems,
+            )
+            results[task_name] = acc
+            print0(f"{task_name} accuracy: {100 * acc:.2f}%")
+
+    # Log to report
+    from nanochat.report import get_report
+    all_tasks_were_evaluated = all(task_name in results for task_name in all_tasks)
+    # calculate the ChatCORE metric if we can (similar to CORE, it's the mean centered accuracy)
+    # this way, ChatCORE ranges from 0 (at random baseline) to 1 (peak performance)
+    chatcore_metric_dict = {}
+    if all_tasks_were_evaluated:
+        centered_mean = 0
+        for task_name, acc in results.items():
+            baseline_acc = baseline_accuracies.get(task_name, 0.0)
+            centered_acc = (acc - baseline_acc) / (1.0 - baseline_acc)
+            centered_mean += centered_acc
+        chatcore_metric = centered_mean / len(results)
+        chatcore_metric_dict = {"ChatCORE metric": chatcore_metric}
+    get_report().log(section="Chat evaluation " + args.source, data=[
+        vars(args), # CLI args
+        results,
+        chatcore_metric_dict,
+    ])
+
+    compute_cleanup()

scripts/chat_rl.py

@@ -0,0 +1,331 @@
+"""
+Reinforcement learning on GSM8K via "GRPO".
+
+I put GRPO in quotes because we actually end up with something a lot
+simpler and more similar to just REINFORCE:
+
+1) Delete trust region, so there is no KL regularization to a reference model
+2) We are on policy, so there's no need for PPO ratio+clip.
+3) We use GAPO style normalization that is token-level, not sequence-level.
+4) Instead of z-score normalization (r - mu)/sigma, only use (r - mu) as the advantage.
+
+1 GPU:
+python -m scripts.chat_rl
+
+8 GPUs:
+torchrun --standalone --nproc_per_node=8 -m scripts.chat_rl -- --run=default
+"""
+
+import os
+import itertools
+import re
+import wandb
+import torch
+import torch.distributed as dist
+
+from nanochat.common import compute_init, compute_cleanup, print0, get_base_dir, DummyWandb
+from nanochat.checkpoint_manager import save_checkpoint, load_model
+from nanochat.engine import Engine
+from tasks.gsm8k import GSM8K
+
+# RL hyperparameters
+run = "dummy" # wandb run name
+source = "sft" # mid|sft
+dtype = "bfloat16"
+device_batch_size = 8 # no forward pass will go above this to not OOM
+examples_per_step = 16 # in total and across all ranks (note: examples, not samples/completions!)
+num_samples = 16 # number of samples per example (/question)
+max_new_tokens = 256
+temperature = 1.0
+top_k = 50 # TODO: try None?
+unembedding_lr = 0.004
+embedding_lr = 0.2
+matrix_lr = 0.02
+weight_decay = 0.0
+init_lr_frac = 0.05
+num_epochs = 1 # how many epochs of gsm8k to train on
+save_every = 60 # every how many steps to save the model
+eval_every = 60 # every how many steps to evaluate the model for val pass@k
+eval_examples = 400 # number of examples used for evaluating pass@k
+# now allow CLI to override the settings via the configurator lol
+config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file
+user_config = {k: globals()[k] for k in config_keys} # will be useful for logging
+# -----------------------------------------------------------------------------
+
+# Init compute/precision
+ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init()
+master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
+dtype = torch.float32 if dtype == 'float32' else torch.bfloat16
+autocast_ctx = torch.amp.autocast(device_type="cuda", dtype=dtype)
+
+# wandb logging init
+use_dummy_wandb = run == "dummy" or not master_process
+wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-rl", name=run, config=user_config)
+
+# Init model and tokenizer
+model, tokenizer, meta = load_model(source, device, phase="eval")
+engine = Engine(model, tokenizer) # for sampling rollouts
+
+# -----------------------------------------------------------------------------
+# Rollout / sampling generator loop that yields batches of examples for training
+
+train_task = GSM8K(subset="main", split="train")
+val_task = GSM8K(subset="main", split="test")
+num_steps = (len(train_task) // examples_per_step) * num_epochs
+print0(f"Calculated number of steps: {num_steps}")
+
+@torch.no_grad()
+def get_batch():
+    assistant_end = tokenizer.encode_special("<|assistant_end|>") # ok to use this token, it's only for padding and isn't used in the loss.
+    rank_indices = range(ddp_rank, len(train_task), ddp_world_size) # each rank is responsible for different examples in the training data
+    for example_idx in itertools.cycle(rank_indices):
+
+        # First get the full conversation of both user and assistant messages
+        conversation = train_task[example_idx]
+
+        # Tokenize the conversation, deleting the last Assistant message and priming the Assistant for a completion instead
+        # (i.e. keep the <|assistant_start|>, but delete everything after it)
+        tokens = tokenizer.render_for_completion(conversation)
+        prefix_length = len(tokens)
+
+        # Generate num_samples samples using batched generation, use loop to avoid OOMs
+        model.eval() # ensure the model is in eval mode
+        generated_token_sequences = []
+        masks = []
+        num_sampling_steps = num_samples // device_batch_size # go sequentially to prevent OOMs
+        for sampling_step in range(num_sampling_steps):
+            seed = hash((step, example_idx, sampling_step)) & 0x7FFFFFFF # positive half of int32
+            with autocast_ctx:
+                generated_token_sequences_batch, masks_batch = engine.generate_batch(
+                    tokens,
+                    num_samples=device_batch_size,
+                    max_tokens=max_new_tokens,
+                    temperature=temperature,
+                    top_k=top_k,
+                    seed=seed, # must make sure to change the seed for each sampling step
+                )
+            generated_token_sequences.extend(generated_token_sequences_batch)
+            masks.extend(masks_batch)
+
+        # Calculate the rewards for each sample
+        rewards = []
+        for sample_tokens in generated_token_sequences:
+            # Get just the generated tokens (after the prompt)
+            generated_tokens = sample_tokens[prefix_length:]
+            # Decode the generated response
+            generated_text = tokenizer.decode(generated_tokens)
+            # Calculate the reward
+            reward = train_task.reward(conversation, generated_text)
+            rewards.append(reward)
+
+        # Pad the sequences so that their lengths (in time) match
+        max_length = max(len(seq) for seq in generated_token_sequences)
+        padded_generated_token_sequences = [seq + [assistant_end] * (max_length - len(seq)) for seq in generated_token_sequences]
+        padded_masks = [mask + [0] * (max_length - len(mask)) for mask in masks]
+        # Stack up the sequences and masks into PyTorch tensors
+        ids = torch.tensor(padded_generated_token_sequences, dtype=torch.long, device=device)
+        mask_ids = torch.tensor(padded_masks, dtype=torch.long, device=device)
+        # Generate autoregressive inputs and targets to the Transformer
+        inputs = ids[:, :-1]
+        targets = ids[:, 1:].clone() # clone to avoid in-place modification:
+        targets[mask_ids[:, 1:] == 0] = -1 # <-- inplace modification right here. -1 is the ignore index
+        # NOTE also that the Engine returns mask=0 for BOTH the prompt tokens AND the tool use tokens.
+        # So we will (correctly) end up not training on the prompt tokens, or the tool use forced tokens.
+        rewards = torch.tensor(rewards, dtype=torch.float, device=device)
+        # Calculate the advantages by simply subtracting the mean (instead of z-score (x-mu)/sigma)
+        mu = rewards.mean()
+        advantages = rewards - mu
+        # yield inputs/targets as (B, T) of ids and rewards as (B,) of floats
+        yield generated_token_sequences, inputs, targets, rewards, advantages
+
+# -----------------------------------------------------------------------------
+# Simple evaluation loop for GSM8K pass@k
+def run_gsm8k_eval(task, tokenizer, engine,
+    max_examples=None,
+    num_samples=1,
+    max_completion_tokens=256,
+    temperature=0.0,
+    top_k=50
+):
+    """
+    Evaluates GSM8K task and returns a list of records of evaluation outcomes.
+    In a distributed setting, all ranks cooperate but this function will NOT
+    do the reduction across ranks. This is the responsibility of the caller.
+    Because the evaluation can take a while, this function will yield records one by one.
+    """
+    max_examples = min(max_examples, len(task)) if max_examples is not None else len(task)
+    for idx in range(ddp_rank, max_examples, ddp_world_size):
+        conversation = task[idx]
+        tokens = tokenizer.render_for_completion(conversation)
+        prefix_length = len(tokens)
+        # Generate k samples using batched generation inside the Engine
+        assert num_samples <= device_batch_size # usually this is true. we can add a loop if not...
+        generated_token_sequences, masks = engine.generate_batch(
+            tokens,
+            num_samples=num_samples,
+            max_tokens=max_completion_tokens,
+            temperature=temperature,
+            top_k=top_k
+        )
+        # Check each sample for correctness
+        outcomes = []
+        for sample_tokens in generated_token_sequences:
+            generated_tokens = sample_tokens[prefix_length:]
+            generated_text = tokenizer.decode(generated_tokens)
+            is_correct = task.evaluate(conversation, generated_text)
+            outcomes.append({
+                "is_correct": is_correct
+            })
+        # A bit bloated because I wanted to do more complex logging at one point.
+        record = {
+            "idx": idx,
+            "outcomes": outcomes,
+        }
+        yield record
+
+# -----------------------------------------------------------------------------
+# Training loop
+
+# Init the optimizer
+optimizers = model.setup_optimizers(
+    unembedding_lr=unembedding_lr,
+    embedding_lr=embedding_lr,
+    matrix_lr=matrix_lr,
+    weight_decay=weight_decay,
+)
+
+# Set the initial learning rate as a fraction of the base learning rate
+for opt in optimizers:
+    for group in opt.param_groups:
+        group["lr"] = group["lr"] * init_lr_frac
+        group["initial_lr"] = group["lr"] # save the initial learning so we can decay easily later
+
+# Learning rate scheduler: simple rampdown to zero over num_steps
+def get_lr_multiplier(it):
+    lrm = 1.0 - it / num_steps
+    return lrm
+
+# Calculate the number of examples each rank handles to achieve the desired examples_per_step
+print0(f"Total sequences per step: {examples_per_step * num_samples}") # total batch size in sequences/step
+assert examples_per_step % ddp_world_size == 0, "Desired examples per step must be divisible by the number of ranks"
+examples_per_rank = examples_per_step // ddp_world_size # per GPU
+print0(f"Calculated examples per rank: {examples_per_rank}")
+
+# Kick off the training loop
+batch_iterator = get_batch()
+for step in range(num_steps):
+
+    # Evaluate the model once in a while and log to wandb
+    if step % eval_every == 0:
+        model.eval()
+        passk = torch.zeros(device_batch_size, device=device) # pass@k for k=1..device_batch_size
+        with autocast_ctx:
+            records_iter = run_gsm8k_eval(val_task, tokenizer, engine, num_samples=device_batch_size, max_examples=eval_examples, temperature=1.0)
+            records = list(records_iter) # collect all records
+        for k in range(1, device_batch_size + 1):
+            passk[k - 1] = sum(any(o["is_correct"] for o in r["outcomes"][:k]) for r in records)
+        num_records = torch.tensor(len(records), dtype=torch.long, device=device)
+        if ddp:
+            dist.all_reduce(num_records, op=dist.ReduceOp.SUM)
+            dist.all_reduce(passk, op=dist.ReduceOp.SUM)
+        passk = passk / num_records.item() # normalize by the total number of records
+        print_passk = [f"Pass@{k}: {passk[k - 1].item():.4f}" for k in range(1, device_batch_size + 1)]
+        print0(f"Step {step} | {', '.join(print_passk)}")
+        log_passk = {f"pass@{k}": passk[k - 1].item() for k in range(1, device_batch_size + 1)}
+        wandb_run.log({
+            "step": step,
+            **log_passk,
+        })
+
+    # Forward/Backward on rollouts over multiple examples in the dataset
+    rewards_list = []
+    sequence_lengths = []
+    for example_step in range(examples_per_rank):
+        # Get one batch corresponding to one example in the training dataset
+        sequences_all, inputs_all, targets_all, rewards_all, advantages_all = next(batch_iterator)
+        # Evaluate the loss and gradients
+        model.train() # ensure the model is in train mode
+        # We need one more loop because we can never exceed the device_batch_size
+        assert inputs_all.size(0) % device_batch_size == 0
+        num_passes = inputs_all.size(0) // device_batch_size
+        for pass_idx in range(num_passes):
+            # Pluck out the batch for this pass
+            b0, b1 = pass_idx * device_batch_size, (pass_idx + 1) * device_batch_size
+            inputs = inputs_all[b0:b1]
+            targets = targets_all[b0:b1]
+            rewards = rewards_all[b0:b1]
+            advantages = advantages_all[b0:b1]
+            # Calculate log probabilities. Note that the loss calculates NLL = -logp, so we negate
+            with autocast_ctx:
+                logp = -model(inputs, targets, loss_reduction='none').view_as(inputs) # (B, T)
+            # Calculate the PG objective. Note that ignore_index=-1 ensures that invalid tokens have loss 0.
+            pg_obj = (logp * advantages.unsqueeze(-1)).sum()
+            # normalize by the number of valid tokens, number of passes, and examples_per_rank
+            num_valid = (targets >= 0).sum().clamp(min=1)
+            pg_obj = pg_obj / (num_valid * num_passes * examples_per_rank)
+            # Note, there is no need to add PPO ratio+clip because we are on policy
+            # Finally, formulate the loss that we want to minimize (instead of objective we wish to maximize)
+            loss = -pg_obj
+            loss.backward()
+            print0(f"Step {step}/{num_steps} | Example step {example_step} | Pass {pass_idx} | loss: {loss.item():.6f} | Average reward: {rewards.mean().item()}")
+        # For logging
+        rewards_list.append(rewards_all.mean().item())
+        sequence_lengths.extend(len(seq) for seq in sequences_all)
+
+    # A bunch of logging for how the rollouts went this step
+    mean_reward = sum(rewards_list) / len(rewards_list)
+    mean_sequence_length = sum(sequence_lengths) / len(sequence_lengths)
+    if ddp: # aggregate across ranks
+        mean_reward_tensor = torch.tensor(mean_reward, dtype=torch.float, device=device)
+        mean_sequence_length_tensor = torch.tensor(mean_sequence_length, dtype=torch.float, device=device)
+        dist.all_reduce(mean_reward_tensor, op=dist.ReduceOp.AVG)
+        dist.all_reduce(mean_sequence_length_tensor, op=dist.ReduceOp.AVG)
+        mean_reward = mean_reward_tensor.item()
+        mean_sequence_length = mean_sequence_length_tensor.item()
+    print0(f"Step {step}/{num_steps} | Average reward: {mean_reward} | Average sequence length: {mean_sequence_length:.2f}")
+    wandb_run.log({
+        "step": step,
+        "reward": mean_reward,
+        "sequence_length": mean_sequence_length,
+    })
+
+    # Update the model parameters
+    lrm = get_lr_multiplier(step)
+    for opt in optimizers: # first set the learning rate
+        for group in opt.param_groups:
+            group["lr"] = group["initial_lr"] * lrm
+    for opt in optimizers: # then step the optimizers
+        opt.step()
+    model.zero_grad(set_to_none=True)
+    wandb_run.log({
+        "step": step,
+        "lrm": lrm,
+    })
+
+    # Master process saves the model once in a while. Skip first step. Save last step.
+    if master_process and ((step > 0 and step % save_every == 0) or step == num_steps - 1):
+        base_dir = get_base_dir()
+        depth = model.config.n_layer
+        model_tag = f"d{depth}" # base the model tag on the depth of the base model
+        checkpoint_dir = os.path.join(base_dir, "chatrl_checkpoints", model_tag)
+        model_config_kwargs = model.config.__dict__ # slightly naughty, abusing the simplicity of GPTConfig, TODO nicer
+        save_checkpoint(
+            checkpoint_dir,
+            step,
+            model.state_dict(),
+            None, # note: we don't bother to save the optimizer state
+            {
+                "model_config": model_config_kwargs,
+            }
+        )
+        print(f"✅ Saved model checkpoint to {checkpoint_dir}")
+
+# Log to report
+from nanochat.report import get_report
+get_report().log(section="Chat RL", data=[
+    user_config, # CLI args
+])
+
+wandb_run.finish() # wandb run finish
+compute_cleanup()

scripts/chat_sft.py

@@ -0,0 +1,285 @@
+"""
+Finetune a base model to be a chat model.
+Run on one GPU e.g. for debugging:
+
+python -m scripts.chat_sft
+
+Or torchrun for training:
+
+torchrun --standalone --nproc_per_node=8 -m scripts.chat_sft
+"""
+
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import wandb
+import torch
+import torch.distributed as dist
+from contextlib import nullcontext
+
+from nanochat.common import compute_init, compute_cleanup, get_base_dir, print0, DummyWandb, autodetect_device_type
+from nanochat.checkpoint_manager import load_model
+from nanochat.checkpoint_manager import save_checkpoint
+from nanochat.engine import Engine
+from scripts.chat_eval import run_chat_eval
+
+from tasks.common import TaskMixture
+from tasks.arc import ARC
+from tasks.gsm8k import GSM8K
+from tasks.smoltalk import SmolTalk
+from tasks.customjson import CustomJSON
+from tasks.spellingbee import SimpleSpelling, SpellingBee
+
+# -----------------------------------------------------------------------------
+# SFT Hyperparameters
+run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb)
+# input model options
+source = "mid" # base|mid , which checkpoint to load the model from (base model or midtrained model)
+model_tag = None # model tag to load the model from (base model or midtrained model)
+step = None # step to load the model from (base model or midtrained model)
+# compute/precision
+device_type = "" # cuda|cpu|mps (empty => autodetect)
+dtype = "bfloat16"
+device_batch_size = 4 # max to avoid OOM
+# optimization
+num_epochs = 1
+num_iterations = -1 # override number of iterations (-1 = disable, use num_epochs to derive it)
+target_examples_per_step = 32
+unembedding_lr = 0.004
+embedding_lr = 0.2
+matrix_lr = 0.02
+weight_decay = 0.0
+init_lr_frac = 0.02
+# evaluation and logging there of
+eval_every = 100
+eval_steps = 100
+eval_metrics_every = 200
+eval_metrics_max_problems = 1024
+# now allow CLI to override the settings via the configurator lol
+config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file
+user_config = {k: globals()[k] for k in config_keys} # possibly useful for logging
+# -----------------------------------------------------------------------------
+
+# Compute init
+device_type = autodetect_device_type() if device_type == "" else device_type
+ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+master_process = ddp_rank == 0
+ptdtype = torch.float32 if dtype == 'float32' else torch.bfloat16
+autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext()
+
+# wandb logging init
+use_dummy_wandb = run == "dummy" or not master_process
+wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-sft", name=run, config=user_config, save_code=True)
+
+# Load the model and tokenizer
+model, tokenizer, meta = load_model(source, device, phase="train", model_tag=model_tag, step=step)
+orig_model = model # original, uncompiled model
+# model = torch.compile(model, dynamic=True) # doesn't work super well because of variable lengths of inputs
+engine = Engine(model, tokenizer) # will be used for inline model evaluation only
+
+# -----------------------------------------------------------------------------
+# Task data mixture we'll train on
+identity_conversations_filepath = os.path.join(get_base_dir(), "identity_conversations.jsonl")
+train_ds = TaskMixture([
+    ARC(subset="ARC-Easy", split="train"), # 2.3K rows
+    ARC(subset="ARC-Challenge", split="train"), # 1.1K rows
+    GSM8K(subset="main", split="train"), # 8K rows
+    SmolTalk(split="train", stop=10_000), # 10K rows of smoltalk
+    CustomJSON(filepath=identity_conversations_filepath), # 1K rows of synthetic identity conversations
+    SimpleSpelling(size=300, split="train"), # 300 rows of Simple Spelling (e.g. spell the word 'apple')
+    SpellingBee(size=300, split="train"), # 300 rows of Spelling Bee (e.g. how many 'r' are in 'strawberry'?)
+]) # 2.3K + 1.1K + 8K + 10K + 1K + 0.3K + 0.3K = 23K rows
+val_ds = SmolTalk(split="test") # general conversations, 24K rows (though we don't actually use all of it)
+
+# -----------------------------------------------------------------------------
+# DataLoader
+
+def sft_data_generator(dataset, batch_size):
+    pad_token_id = tokenizer.encode_special("<|assistant_end|>") # use <|assistant_end|> as the pad token is ok, these positions are masked in the loss
+    # prepares a list of tokenized conversations into a batch and yields
+    def collate_and_yield(batch):
+        nrows = len(batch)
+        ncols = max(len(ids) for ids, mask in batch) - 1 # seq of n creates inputs/targets of n-1
+        inputs = torch.full((nrows, ncols), pad_token_id, dtype=torch.long)
+        targets = torch.full((nrows, ncols), -1, dtype=torch.long) # -1 is ignore index
+        for i, (ids, mask) in enumerate(batch):
+            n = len(ids)
+            ids_tensor = torch.tensor(ids, dtype=torch.long)
+            inputs[i, :n-1] = ids_tensor[:-1]
+            # recall -1 is the ignore index, so mask out targets where mask is 0
+            row_targets = ids_tensor[1:]
+            # mask[1:] omits the mask for the BOS token, which is never a target atm so it's ok
+            mask_tensor = torch.tensor(mask[1:], dtype=torch.long)
+            row_targets[mask_tensor == 0] = -1 # mask out targets where mask is 0
+            targets[i, :n-1] = row_targets
+        inputs = inputs.to(device) # move to device
+        targets = targets.to(device)
+        return inputs, targets
+    # iterates over the dataset in epochs, tokenizes
+    batch = []
+    while True:
+        for i in range(ddp_rank, len(dataset), ddp_world_size):
+            doc = dataset[i]
+            ids, mask = tokenizer.render_conversation(doc)
+            batch.append((ids, mask))
+            if len(batch) == batch_size:
+                yield collate_and_yield(batch)
+                batch = []
+
+examples_per_step = device_batch_size * ddp_world_size
+print0(f"Target examples per step: {target_examples_per_step}")
+print0(f"Device batch size: {device_batch_size}")
+print0(f"Examples per step is device_batch_size * ddp_world_size: {examples_per_step}")
+assert target_examples_per_step % examples_per_step == 0, "Target examples per step must be divisible by examples per step"
+grad_accum_steps = target_examples_per_step // examples_per_step
+print0(f"=> Setting grad accum steps: {grad_accum_steps}")
+
+if num_iterations == -1:
+    # derive num_iterations from num_epochs and the size of the dataset
+    assert num_epochs > 0, "num_epochs must be positive if num_iterations is -1"
+    num_iterations = (len(train_ds) // target_examples_per_step) * num_epochs
+train_loader = sft_data_generator(train_ds, batch_size=device_batch_size)
+build_val_loader = lambda: sft_data_generator(val_ds, batch_size=device_batch_size)
+
+# -----------------------------------------------------------------------------
+# Initialize the Optimizer
+
+optimizers = model.setup_optimizers(
+    unembedding_lr=unembedding_lr,
+    embedding_lr=embedding_lr,
+    matrix_lr=matrix_lr,
+    weight_decay=weight_decay,
+)
+# Set the initial learning rate as a fraction of the base learning rate
+for opt in optimizers:
+    for group in opt.param_groups:
+        group["lr"] = group["lr"] * init_lr_frac
+        group["initial_lr"] = group["lr"] # save the initial learning so we can decay easily later
+
+# -----------------------------------------------------------------------------
+# Training loop
+
+# Learning rate scheduler
+def get_lr_multiplier(it):
+    lrm = 1.0 - it / num_iterations
+    return lrm
+
+# Go!
+step = 0
+train_iter = iter(train_loader)
+for step in range(num_iterations):
+    last_step = step == num_iterations - 1
+
+    # evaluate the validation loss
+    if last_step or step % eval_every == 0:
+        model.eval()
+        val_iter = iter(build_val_loader())
+        losses = []
+        for _ in range(eval_steps):
+            val_inputs, val_targets = next(val_iter)
+            with torch.no_grad(), autocast_ctx:
+                loss = model(val_inputs, val_targets)
+            losses.append(loss)
+        val_loss = torch.stack(losses).mean() # average over eval_steps
+        if ddp:
+            dist.all_reduce(val_loss, op=dist.ReduceOp.AVG) # average over ranks
+        val_loss = val_loss.item()
+        print0(f"Step {step:05d} | Validation loss: {val_loss:.6f}")
+        wandb_run.log({
+            "step": step,
+            "val_loss": val_loss,
+        })
+        model.train()
+
+    # evlauate accuracy of the multiple choice tasks (which are quick to run)
+    if last_step or (step > 0 and step % eval_metrics_every == 0):
+        model.eval()
+        metrics = {}
+        with torch.no_grad(), autocast_ctx:
+            # note that because these are inside no_grad, we can usually afford to at least ~2X the batch size
+            metrics["mmlu_acc"] = run_chat_eval("MMLU", model, tokenizer, engine, batch_size=device_batch_size*2, max_problems=eval_metrics_max_problems)
+            metrics["arc_easy_acc"] = run_chat_eval("ARC-Easy", model, tokenizer, engine, batch_size=device_batch_size*2, max_problems=eval_metrics_max_problems)
+        metrics_str = ', '.join(f'{k}: {v:.6f}' for k, v in metrics.items())
+        print0(f"Step {step:05d} | {metrics_str}")
+        wandb_run.log({
+            "step": step,
+            **metrics,
+        })
+        model.train()
+
+    if last_step:
+        break
+
+    # evaluate the gradient
+    num_tokens = torch.tensor(0, device=device) # the number of "active" tokens of supervision seen
+    for micro_step in range(grad_accum_steps):
+        train_inputs, train_targets = next(train_iter)
+        with autocast_ctx:
+            loss = model(train_inputs, train_targets)
+        train_loss = loss.detach() # for logging
+        loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here
+        loss.backward() # accumulate the gradient
+        num_tokens += (train_targets >= 0).sum()
+    if ddp:
+        dist.all_reduce(num_tokens, op=dist.ReduceOp.SUM) # sum over ranks
+
+    # learning rate scheduler
+    lrm = get_lr_multiplier(step)
+    for opt in optimizers:
+        for group in opt.param_groups:
+            group["lr"] = group["initial_lr"] * lrm
+
+    # step the optimizers
+    for opt in optimizers:
+        opt.step()
+    model.zero_grad(set_to_none=True)
+
+    # logging
+    train_loss_item = train_loss.item()
+    num_tokens_item = num_tokens.item()
+    print0(f"Step {step:05d}/{num_iterations:05d} | Training loss: {train_loss_item:.6f}| lrm: {lrm:.6f}| num_tokens: {num_tokens_item:,}")
+    wandb_run.log({
+        "step": step,
+        "lrm": lrm,
+        "train_loss": train_loss_item,
+        "num_tokens": num_tokens_item,
+    })
+    step += 1
+
+# Save the model at the end of the run
+if master_process:
+    base_dir = get_base_dir()
+    depth = model.config.n_layer
+    model_tag = f"d{depth}" # base the model tag on the depth of the base model
+    checkpoint_dir = os.path.join(base_dir, "chatsft_checkpoints", model_tag)
+    model_config_kwargs = model.config.__dict__ # slightly naughty, abusing the simplicity of GPTConfig, TODO nicer
+    save_checkpoint(
+        checkpoint_dir,
+        step,
+        model.state_dict(),
+        None, # note: we don't bother to save the optimizer state
+        {
+            "step": step,
+            "val_loss": val_loss,
+            **metrics,
+            "model_config": model_config_kwargs,
+        }
+    )
+    print(f"✅ Saved model checkpoint to {checkpoint_dir}")
+
+# Log to report
+from nanochat.report import get_report
+get_report().log(section="Chat SFT", data=[
+    user_config, # CLI args
+    {
+        "Training rows": len(train_ds),
+        "Number of iterations": num_iterations,
+        "Training loss": train_loss_item,
+        "Validation loss": val_loss,
+    },
+])
+
+# Cleanup
+wandb_run.finish()
+compute_cleanup()

scripts/chat_web.py

@@ -0,0 +1,415 @@
+#!/usr/bin/env python3
+"""
+Unified web chat server - serves both UI and API from a single FastAPI instance.
+
+Uses data parallelism to distribute requests across multiple GPUs. Each GPU loads
+a full copy of the model, and incoming requests are distributed to available workers.
+
+Launch examples:
+
+- single available GPU (default)
+python -m scripts.chat_web
+
+- 4 GPUs
+python -m scripts.chat_web --num-gpus 4
+
+To chat, open the URL printed in the console. (If on cloud box, make sure to use public IP)
+
+Endpoints:
+  GET  /           - Chat UI
+  POST /chat/completions - Chat API (streaming only)
+  GET  /health     - Health check with worker pool status
+  GET  /stats      - Worker pool statistics and GPU utilization
+
+Abuse Prevention:
+  - Maximum 500 messages per request
+  - Maximum 8000 characters per message
+  - Maximum 32000 characters total conversation length
+  - Temperature clamped to 0.0-2.0
+  - Top-k clamped to 1-200
+  - Max tokens clamped to 1-4096
+"""
+
+import argparse
+import json
+import os
+import torch
+import asyncio
+import logging
+import random
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, HTMLResponse, FileResponse
+from pydantic import BaseModel
+from typing import List, Optional, AsyncGenerator
+from dataclasses import dataclass
+from contextlib import nullcontext
+from nanochat.common import compute_init, autodetect_device_type
+from nanochat.checkpoint_manager import load_model
+from nanochat.engine import Engine
+
+# Abuse prevention limits
+MAX_MESSAGES_PER_REQUEST = 500
+MAX_MESSAGE_LENGTH = 8000
+MAX_TOTAL_CONVERSATION_LENGTH = 32000
+MIN_TEMPERATURE = 0.0
+MAX_TEMPERATURE = 2.0
+MIN_TOP_K = 1
+MAX_TOP_K = 200
+MIN_MAX_TOKENS = 1
+MAX_MAX_TOKENS = 4096
+
+parser = argparse.ArgumentParser(description='NanoChat Web Server')
+parser.add_argument('-n', '--num-gpus', type=int, default=1, help='Number of GPUs to use (default: 1)')
+parser.add_argument('-i', '--source', type=str, default="sft", help="Source of the model: sft|mid|rl")
+parser.add_argument('-t', '--temperature', type=float, default=0.8, help='Default temperature for generation')
+parser.add_argument('-k', '--top-k', type=int, default=50, help='Default top-k sampling parameter')
+parser.add_argument('-m', '--max-tokens', type=int, default=512, help='Default max tokens for generation')
+parser.add_argument('-g', '--model-tag', type=str, default=None, help='Model tag to load')
+parser.add_argument('-s', '--step', type=int, default=None, help='Step to load')
+parser.add_argument('-p', '--port', type=int, default=8000, help='Port to run the server on')
+parser.add_argument('-d', '--dtype', type=str, default='bfloat16', choices=['float32', 'bfloat16'])
+parser.add_argument('--device-type', type=str, default='', choices=['cuda', 'cpu', 'mps'], help='Device type for evaluation: cuda|cpu|mps. empty => autodetect')
+parser.add_argument('--host', type=str, default='0.0.0.0', help='Host to bind the server to')
+args = parser.parse_args()
+
+# Configure logging for conversation traffic
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+device_type = autodetect_device_type() if args.device_type == "" else args.device_type
+ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+ptdtype = torch.float32 if args.dtype == 'float32' else torch.bfloat16
+
+@dataclass
+class Worker:
+    """A worker with a model loaded on a specific GPU."""
+    gpu_id: int
+    device: torch.device
+    engine: Engine
+    tokenizer: object
+    autocast_ctx: torch.amp.autocast
+
+class WorkerPool:
+    """Pool of workers, each with a model replica on a different GPU."""
+
+    def __init__(self, num_gpus: Optional[int] = None):
+        if num_gpus is None:
+            if device_type == "cuda":
+                num_gpus = torch.cuda.device_count()
+            else:
+                num_gpus = 1 # e.g. cpu|mps
+        self.num_gpus = num_gpus
+        self.workers: List[Worker] = []
+        self.available_workers: asyncio.Queue = asyncio.Queue()
+
+    async def initialize(self, source: str, model_tag: Optional[str] = None, step: Optional[int] = None):
+        """Load model on each GPU."""
+        print(f"Initializing worker pool with {self.num_gpus} GPUs...")
+        if self.num_gpus > 1:
+            assert device_type == "cuda", "Only CUDA supports multiple workers/GPUs. cpu|mps does not."
+
+        for gpu_id in range(self.num_gpus):
+
+            if device_type == "cuda":
+                device = torch.device(f"cuda:{gpu_id}")
+                print(f"Loading model on GPU {gpu_id}...")
+            else:
+                device = torch.device(device_type) # e.g. cpu|mps
+                print(f"Loading model on {device_type}...")
+
+            model, tokenizer, _ = load_model(source, device, phase="eval", model_tag=model_tag, step=step)
+            engine = Engine(model, tokenizer)
+            autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if device_type == "cuda" else nullcontext()
+
+            worker = Worker(
+                gpu_id=gpu_id,
+                device=device,
+                engine=engine,
+                tokenizer=tokenizer,
+                autocast_ctx=autocast_ctx
+            )
+            self.workers.append(worker)
+            await self.available_workers.put(worker)
+
+        print(f"All {self.num_gpus} workers initialized!")
+
+    async def acquire_worker(self) -> Worker:
+        """Get an available worker from the pool."""
+        return await self.available_workers.get()
+
+    async def release_worker(self, worker: Worker):
+        """Return a worker to the pool."""
+        await self.available_workers.put(worker)
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatRequest(BaseModel):
+    messages: List[ChatMessage]
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+    top_k: Optional[int] = None
+
+def validate_chat_request(request: ChatRequest):
+    """Validate chat request to prevent abuse."""
+    # Check number of messages
+    if len(request.messages) == 0:
+        raise HTTPException(status_code=400, detail="At least one message is required")
+    if len(request.messages) > MAX_MESSAGES_PER_REQUEST:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Too many messages. Maximum {MAX_MESSAGES_PER_REQUEST} messages allowed per request"
+        )
+
+    # Check individual message lengths and total conversation length
+    total_length = 0
+    for i, message in enumerate(request.messages):
+        if not message.content:
+            raise HTTPException(status_code=400, detail=f"Message {i} has empty content")
+
+        msg_length = len(message.content)
+        if msg_length > MAX_MESSAGE_LENGTH:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Message {i} is too long. Maximum {MAX_MESSAGE_LENGTH} characters allowed per message"
+            )
+        total_length += msg_length
+
+    if total_length > MAX_TOTAL_CONVERSATION_LENGTH:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Total conversation is too long. Maximum {MAX_TOTAL_CONVERSATION_LENGTH} characters allowed"
+        )
+
+    # Validate role values
+    for i, message in enumerate(request.messages):
+        if message.role not in ["user", "assistant"]:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Message {i} has invalid role. Must be 'user', 'assistant', or 'system'"
+            )
+
+    # Validate temperature
+    if request.temperature is not None:
+        if not (MIN_TEMPERATURE <= request.temperature <= MAX_TEMPERATURE):
+            raise HTTPException(
+                status_code=400,
+                detail=f"Temperature must be between {MIN_TEMPERATURE} and {MAX_TEMPERATURE}"
+            )
+
+    # Validate top_k
+    if request.top_k is not None:
+        if not (MIN_TOP_K <= request.top_k <= MAX_TOP_K):
+            raise HTTPException(
+                status_code=400,
+                detail=f"top_k must be between {MIN_TOP_K} and {MAX_TOP_K}"
+            )
+
+    # Validate max_tokens
+    if request.max_tokens is not None:
+        if not (MIN_MAX_TOKENS <= request.max_tokens <= MAX_MAX_TOKENS):
+            raise HTTPException(
+                status_code=400,
+                detail=f"max_tokens must be between {MIN_MAX_TOKENS} and {MAX_MAX_TOKENS}"
+            )
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Load models on all GPUs on startup."""
+    print("Loading nanochat models across GPUs...")
+    app.state.worker_pool = WorkerPool(num_gpus=args.num_gpus)
+    await app.state.worker_pool.initialize(args.source, model_tag=args.model_tag, step=args.step)
+    print(f"Server ready at http://localhost:{args.port}")
+    yield
+
+app = FastAPI(lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/")
+async def root():
+    """Serve the chat UI."""
+    ui_html_path = os.path.join("nanochat", "ui.html")
+    with open(ui_html_path, "r", encoding="utf-8") as f:
+        html_content = f.read()
+    # Replace the API_URL to use the same origin
+    html_content = html_content.replace(
+        "const API_URL = `http://${window.location.hostname}:8000`;",
+        "const API_URL = '';"
+    )
+    return HTMLResponse(content=html_content)
+
+
+@app.get("/logo.svg")
+async def logo():
+    """Serve the NanoChat logo for favicon and header."""
+    logo_path = os.path.join("nanochat", "logo.svg")
+    return FileResponse(logo_path, media_type="image/svg+xml")
+
+async def generate_stream(
+    worker: Worker,
+    tokens,
+    temperature=None,
+    max_new_tokens=None,
+    top_k=None
+) -> AsyncGenerator[str, None]:
+    """Generate assistant response with streaming."""
+    temperature = temperature if temperature is not None else args.temperature
+    max_new_tokens = max_new_tokens if max_new_tokens is not None else args.max_tokens
+    top_k = top_k if top_k is not None else args.top_k
+
+    assistant_end = worker.tokenizer.encode_special("<|assistant_end|>")
+    bos = worker.tokenizer.get_bos_token_id()
+
+    # Accumulate tokens to properly handle multi-byte UTF-8 characters (like emojis)
+    accumulated_tokens = []
+    # Track the last complete UTF-8 string (without replacement characters)
+    last_clean_text = ""
+
+    with worker.autocast_ctx:
+        for token_column, token_masks in worker.engine.generate(
+            tokens,
+            num_samples=1,
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            seed=random.randint(0, 2**31 - 1)
+        ):
+            token = token_column[0]
+
+            # Stopping criteria
+            if token == assistant_end or token == bos:
+                break
+
+            # Append the token to sequence
+            accumulated_tokens.append(token)
+            # Decode all accumulated tokens to get proper UTF-8 handling
+            # Note that decode is a quite efficient operation, basically table lookup and string concat
+            current_text = worker.tokenizer.decode(accumulated_tokens)
+            # Only emit text if it doesn't end with a replacement character
+            # This ensures we don't emit incomplete UTF-8 sequences
+            if not current_text.endswith('�'):
+                # Extract only the new text since last clean decode
+                new_text = current_text[len(last_clean_text):]
+                if new_text:  # Only yield if there's new content
+                    yield f"data: {json.dumps({'token': new_text, 'gpu': worker.gpu_id}, ensure_ascii=False)}\n\n"
+                    last_clean_text = current_text
+
+    yield f"data: {json.dumps({'done': True})}\n\n"
+
+@app.post("/chat/completions")
+async def chat_completions(request: ChatRequest):
+    """Chat completion endpoint (streaming only) - uses worker pool for multi-GPU."""
+
+    # Basic validation to prevent abuse
+    validate_chat_request(request)
+
+    # Log incoming conversation to console
+    logger.info("="*20)
+    for i, message in enumerate(request.messages):
+        logger.info(f"[{message.role.upper()}]: {message.content}")
+    logger.info("-"*20)
+
+    # Acquire a worker from the pool (will wait if all are busy)
+    worker_pool = app.state.worker_pool
+    worker = await worker_pool.acquire_worker()
+
+    try:
+        # Build conversation tokens
+        bos = worker.tokenizer.get_bos_token_id()
+        user_start = worker.tokenizer.encode_special("<|user_start|>")
+        user_end = worker.tokenizer.encode_special("<|user_end|>")
+        assistant_start = worker.tokenizer.encode_special("<|assistant_start|>")
+        assistant_end = worker.tokenizer.encode_special("<|assistant_end|>")
+
+        conversation_tokens = [bos]
+        for message in request.messages:
+            if message.role == "user":
+                conversation_tokens.append(user_start)
+                conversation_tokens.extend(worker.tokenizer.encode(message.content))
+                conversation_tokens.append(user_end)
+            elif message.role == "assistant":
+                conversation_tokens.append(assistant_start)
+                conversation_tokens.extend(worker.tokenizer.encode(message.content))
+                conversation_tokens.append(assistant_end)
+
+        conversation_tokens.append(assistant_start)
+
+        # Streaming response with worker release after completion
+        response_tokens = []
+        async def stream_and_release():
+            try:
+                async for chunk in generate_stream(
+                    worker,
+                    conversation_tokens,
+                    temperature=request.temperature,
+                    max_new_tokens=request.max_tokens,
+                    top_k=request.top_k
+                ):
+                    # Accumulate response for logging
+                    chunk_data = json.loads(chunk.replace("data: ", "").strip())
+                    if "token" in chunk_data:
+                        response_tokens.append(chunk_data["token"])
+                    yield chunk
+            finally:
+                # Log the assistant response to console
+                full_response = "".join(response_tokens)
+                logger.info(f"[ASSISTANT] (GPU {worker.gpu_id}): {full_response}")
+                logger.info("="*20)
+                # Release worker back to pool after streaming is done
+                await worker_pool.release_worker(worker)
+
+        return StreamingResponse(
+            stream_and_release(),
+            media_type="text/event-stream"
+        )
+    except Exception as e:
+        # Make sure to release worker even on error
+        await worker_pool.release_worker(worker)
+        raise e
+
+@app.get("/health")
+async def health():
+    """Health check endpoint."""
+    worker_pool = getattr(app.state, 'worker_pool', None)
+    return {
+        "status": "ok",
+        "ready": worker_pool is not None and len(worker_pool.workers) > 0,
+        "num_gpus": worker_pool.num_gpus if worker_pool else 0,
+        "available_workers": worker_pool.available_workers.qsize() if worker_pool else 0
+    }
+
+@app.get("/stats")
+async def stats():
+    """Get worker pool statistics."""
+    worker_pool = app.state.worker_pool
+    return {
+        "total_workers": len(worker_pool.workers),
+        "available_workers": worker_pool.available_workers.qsize(),
+        "busy_workers": len(worker_pool.workers) - worker_pool.available_workers.qsize(),
+        "workers": [
+            {
+                "gpu_id": w.gpu_id,
+                "device": str(w.device)
+            } for w in worker_pool.workers
+        ]
+    }
+
+if __name__ == "__main__":
+    import uvicorn
+    print(f"Starting NanoChat Web Server")
+    print(f"Temperature: {args.temperature}, Top-k: {args.top_k}, Max tokens: {args.max_tokens}")
+    uvicorn.run(app, host=args.host, port=args.port)

scripts/mid_train.py

@@ -0,0 +1,311 @@
+"""
+Midtrain the model. Same as pretraining but simpler.
+Run as:
+
+python -m scripts.mid_train
+
+Or torchrun for training:
+
+torchrun --standalone --nproc_per_node=8 -m scripts.mid_train -- --device_batch_size=16
+"""
+
+from collections import deque
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+import time
+import wandb
+import torch
+from contextlib import nullcontext
+from nanochat.common import compute_init, compute_cleanup, print0, DummyWandb, get_base_dir, autodetect_device_type
+from nanochat.tokenizer import get_token_bytes
+from nanochat.checkpoint_manager import save_checkpoint
+from nanochat.loss_eval import evaluate_bpb
+from nanochat.checkpoint_manager import load_model
+import torch.distributed as dist
+
+from tasks.common import TaskMixture
+from tasks.gsm8k import GSM8K
+from tasks.mmlu import MMLU
+from tasks.smoltalk import SmolTalk
+from tasks.customjson import CustomJSON
+from tasks.spellingbee import SimpleSpelling, SpellingBee
+
+# -----------------------------------------------------------------------------
+run = "dummy" # wandb run name default ("dummy" is special - we won't log to wandb)
+device_type = "" # cuda|cpu|mps (empty => autodetect)
+model_tag = None # model tag to load the model from (base model or midtrained model)
+step = None # step to load the model from (base model or midtrained model)
+dtype = "bfloat16"
+num_iterations = -1 # explicit number of steps of the optimization (-1 = disable)
+max_seq_len = 2048
+device_batch_size = 32
+unembedding_lr = 0.004
+embedding_lr = 0.2
+matrix_lr = 0.02
+init_lr_frac = 1.0 # initial learning rate is this fraction of the base learning rate
+weight_decay = 0.0
+eval_every = 150 # -1 = disable
+eval_tokens = 20*524288
+total_batch_size = 524288
+dry_run = 0 # dry_run=1 is for experiments: we will log to wandb but we won't write checkpoints or report
+config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
+exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file
+user_config = {k: globals()[k] for k in config_keys} # possibly useful for logging
+# -----------------------------------------------------------------------------
+
+# Compute init
+device_type = autodetect_device_type() if device_type == "" else device_type
+ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type)
+master_process = ddp_rank == 0
+autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
+synchronize = torch.cuda.synchronize if device_type == "cuda" else lambda: None
+get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else lambda: 0
+
+# wandb logging init
+use_dummy_wandb = run == "dummy" or not master_process
+wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-mid", name=run, config=user_config)
+
+# Load the model and tokenizer
+model, tokenizer, meta = load_model("base", device, phase="train", model_tag=model_tag, step=step)
+pretrain_batch_size = meta.get("device_batch_size", None)
+if pretrain_batch_size is not None and device_batch_size > pretrain_batch_size:
+    print0(f"FOOTGUN WARNING: base model training used device_batch_size {pretrain_batch_size}, did you pass in a good --device_batch_size to this script?")
+orig_model = model
+model = torch.compile(model, dynamic=False)
+depth = model.config.n_layer
+num_flops_per_token = model.estimate_flops()
+tokens_per_fwdbwd = device_batch_size * max_seq_len # tokens per iteration for a single rank
+world_tokens_per_fwdbwd = tokens_per_fwdbwd * ddp_world_size # total tokens per iteration for all ranks
+assert total_batch_size % world_tokens_per_fwdbwd == 0
+grad_accum_steps = total_batch_size // world_tokens_per_fwdbwd
+print0(f"Tokens / micro-batch / rank: {device_batch_size} x {max_seq_len} = {tokens_per_fwdbwd:,}")
+print0(f"Tokens / micro-batch: {world_tokens_per_fwdbwd:,}")
+print0(f"Total batch size {total_batch_size:,} => gradient accumulation steps: {grad_accum_steps}")
+token_bytes = get_token_bytes(device=device)
+
+# Initialize the Optimizer (Muon for Linear layers, AdamW for embedding and lm_head)
+optimizers = model.setup_optimizers(unembedding_lr=unembedding_lr, embedding_lr=embedding_lr, matrix_lr=matrix_lr, weight_decay=weight_decay)
+adamw_optimizer, muon_optimizer = optimizers
+# Override the initial learning rate as a fraction of the base learning rate
+for opt in optimizers:
+    for group in opt.param_groups:
+        group["lr"] = group["lr"] * init_lr_frac
+        group["initial_lr"] = group["lr"] # save the initial learning so we can decay easily later
+
+# Midtraining data mixture and DataLoader
+base_dir = get_base_dir()
+identity_conversations_filepath = os.path.join(base_dir, "identity_conversations.jsonl")
+train_dataset = TaskMixture([
+    SmolTalk(split="train"), # 460K rows of general conversations
+    MMLU(subset="auxiliary_train", split="train"), # 100K rows of multiple choice problems drawn from ARC, MC_TEST, OBQA, RACE
+    GSM8K(subset="main", split="train"), # 8K rows teaching simple math and (calculator) tool use
+    CustomJSON(filepath=identity_conversations_filepath), # 1000 rows of synthetic identity conversations
+    CustomJSON(filepath=identity_conversations_filepath), # let's do 2 epochs of these
+    SimpleSpelling(size=200000, split="train"), # 200K rows of Simple Spelling (e.g. spell the word 'apple')
+    SpellingBee(size=80000, split="train"), # 80K rows of Spelling Bee (e.g. how many 'r' are in 'strawberry'?)
+]) # total: 460K + 100K + 8K + 200K + 80K = 848K rows
+val_dataset = TaskMixture([
+    SmolTalk(split="test"), # 24K rows in test set
+    MMLU(subset="all", split="test", stop=5200), # 14K rows in test set, use only 5.2K to match the train ratios
+    GSM8K(subset="main", split="test", stop=420), # 1.32K rows in test set, use only 420 to match the train ratios
+]) # total: 24K + 14K + 1.32K ~= 39K rows
+# DataLoader is defined here, it emits inputs, targets : 2D tensors of shape (device_batch_size, max_seq_len)
+# A big problem is that we don't know the final num_iterations in advance. So we create
+# these two global variables and update them from within the data generator.
+last_step = False # we will toggle this to True when we reach the end of the dataset
+approx_progress = 0.0 # will go from 0 to 1 over the course of the epoch
+def mid_data_generator(split):
+    global last_step, approx_progress
+    assert split in {"train", "val"}, "split must be 'train' or 'val'"
+    dataset = train_dataset if split == "train" else val_dataset
+    dataset_size = len(dataset)
+    assert dataset_size > 0
+    needed_tokens = device_batch_size * max_seq_len + 1 # to form one training batch of inputs,targets
+    token_buffer = deque()
+    # CUDA supports memory pinning for faster transfers between CPU and GPU:
+    scratch = torch.empty(needed_tokens, dtype=torch.int64, pin_memory=(device_type == "cuda"))
+    cursor = ddp_rank # increments by ddp_world_size each time, so each rank processes unique documents
+    it = 0 # iteration counter
+    while True:
+        # Accumulate enough tokens for one iteration before yielding
+        while len(token_buffer) < needed_tokens:
+            conversation = dataset[cursor]
+            ids, _ = tokenizer.render_conversation(conversation)
+            token_buffer.extend(ids)
+            cursor += ddp_world_size
+            if cursor >= dataset_size:
+                cursor -= dataset_size # wrap around for another epoch
+                if split == "train":
+                    last_step = True # toggle last_step to True, which will terminate the training loop
+        # Stopping condition to respect num_iterations, if given
+        it += 1
+        if num_iterations > 0 and it >= num_iterations:
+            last_step = True # toggle last_step to True, which will terminate the training loop
+        # Build up inputs/targets and yield
+        for i in range(needed_tokens):
+            scratch[i] = token_buffer.popleft()
+        inputs_cpu = scratch[:-1].to(dtype=torch.int32)
+        targets_cpu = scratch[1:]
+        inputs = inputs_cpu.view(device_batch_size, max_seq_len).to(device=device, dtype=torch.int32, non_blocking=True)
+        targets = targets_cpu.view(device_batch_size, max_seq_len).to(device=device, dtype=torch.int64, non_blocking=True)
+        if split == "train":
+            if num_iterations > 0:
+                approx_progress = it / num_iterations # calculate progress from the max number of iterations
+            else:
+                approx_progress = cursor / dataset_size # approximate progress as a fraction of the dataset
+        yield inputs, targets
+
+train_loader = mid_data_generator("train")
+build_val_loader = lambda: mid_data_generator("val")
+progress = 0 # will go from 0 to 1 over the course of the epoch
+
+# Learning rate scheduler
+def get_lr_multiplier(progress):
+    # first 80% of training: no decay, then linearly ramp down to 0.
+    return 1 if progress < 0.8 else 1 - (progress - 0.8) / 0.2
+
+# Momentum scheduler for Muon optimizer
+def get_muon_momentum(it):
+    frac = min(it / 300, 1)
+    momentum = (1 - frac) * 0.85 + frac * 0.95
+    return momentum
+
+# -----------------------------------------------------------------------------
+# Training loop
+x, y = next(train_loader) # prefetch the very first batch of data
+min_val_bpb = float("inf")
+smooth_train_loss = 0 # EMA of training loss
+ema_beta = 0.9 # EMA decay factor
+total_training_time = 0 # total wall-clock time of training
+step = 0
+while True:
+    flops_so_far = num_flops_per_token * total_batch_size * step
+
+    # Synchronize last_step across all ranks to avoid hangs in the distributed setting
+    if ddp:
+        last_step_tensor = torch.tensor(last_step, dtype=torch.int32, device=device)
+        dist.all_reduce(last_step_tensor, op=dist.ReduceOp.MAX)
+        last_step = bool(last_step_tensor.item())
+
+    # once in a while: evaluate the val bpb (all ranks participate)
+    if eval_every > 0 and (last_step or step % eval_every == 0):
+        model.eval()
+        val_loader = build_val_loader()
+        eval_steps = eval_tokens // (device_batch_size * max_seq_len * ddp_world_size)
+        with autocast_ctx:
+            val_bpb = evaluate_bpb(model, val_loader, eval_steps, token_bytes)
+        print0(f"Step {step:05d} | Validation bpb: {val_bpb:.4f}")
+        if val_bpb < min_val_bpb:
+            min_val_bpb = val_bpb
+        wandb_run.log({
+            "step": step,
+            "total_training_flops": flops_so_far,
+            "total_training_time": total_training_time,
+            "val/bpb": val_bpb,
+        })
+        model.train()
+
+    # save checkpoint at the end of the run (only on master process)
+    if master_process and last_step and not dry_run:
+        output_dirname = f"d{depth}" # e.g. d12
+        checkpoint_dir = os.path.join(base_dir, "mid_checkpoints", output_dirname)
+        save_checkpoint(
+            checkpoint_dir,
+            step,
+            orig_model.state_dict(),
+            [opt.state_dict() for opt in optimizers], # TODO: make sure saving across ranks is done correctly
+            {
+                "step": step,
+                "val_bpb": val_bpb, # loss at last step
+                "model_config": {
+                    "sequence_len": max_seq_len,
+                    "vocab_size": tokenizer.get_vocab_size(),
+                    "n_layer": depth,
+                    "n_head": model.config.n_head,
+                    "n_kv_head": model.config.n_kv_head,
+                    "n_embd": model.config.n_embd,
+                },
+                "user_config": user_config, # inputs to the training script
+            }
+        )
+
+    if last_step:
+        break
+
+    # -------------------------------------------------------------------------
+    # single training step
+    # evaluate the gradient
+    synchronize()
+    t0 = time.time()
+    for micro_step in range(grad_accum_steps):
+        with autocast_ctx:
+            loss = model(x, y)
+        train_loss = loss.detach() # for logging
+        loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here
+        loss.backward()
+        x, y = next(train_loader) # prefetch the next batch while the GPU is busy with forward/backward
+        progress = max(progress, approx_progress) # only increase progress monotonically
+    # step the optimizers
+    lrm = get_lr_multiplier(progress)
+    for opt in optimizers:
+        for group in opt.param_groups:
+            group["lr"] = group["initial_lr"] * lrm
+    muon_momentum = get_muon_momentum(step)
+    for group in muon_optimizer.param_groups:
+        group["momentum"] = muon_momentum
+    for opt in optimizers:
+        opt.step()
+    model.zero_grad(set_to_none=True)
+    synchronize()
+    t1 = time.time()
+    dt = t1 - t0
+    # -------------------------------------------------------------------------
+
+    # State
+    step += 1
+
+    # logging
+    smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss.item() # EMA the training loss
+    debiased_smooth_loss = smooth_train_loss / (1 - ema_beta**(step + 1)) # debias the EMA
+    pct_done = 100 * progress
+    tok_per_sec = int(total_batch_size / dt)
+    flops_per_sec = num_flops_per_token * total_batch_size / dt
+    promised_flops_per_sec_h100 = 989e12 * ddp_world_size # bfloat16 H100 SXM and without 2:4 sparsity
+    mfu = 100 * flops_per_sec / promised_flops_per_sec_h100 # in %
+    if step > 10:
+        total_training_time += dt # only count the time after the first 10 steps
+    print0(f"step {step:05d} ({pct_done:.2f}%) | loss: {debiased_smooth_loss:.6f} | lrm: {lrm:.2f} | dt: {dt * 1000:.2f}ms | tok/sec: {tok_per_sec:,} | mfu: {mfu:.2f} | total time: {total_training_time/60:.2f}m")
+    if step % 10 == 0:
+        wandb_run.log({
+            "step": step,
+            "total_training_flops": flops_so_far,
+            "total_training_time": total_training_time,
+            "train/loss": debiased_smooth_loss,
+            "train/lrm": lrm,
+            "train/dt": dt,
+            "train/tok_per_sec": tok_per_sec,
+            "train/mfu": mfu,
+        })
+
+# print a few more stats
+print0(f"Peak memory usage: {get_max_memory() / 1024 / 1024:.2f}MiB")
+print0(f"Total training time: {total_training_time/60:.2f}m")
+print0(f"Minimum validation bpb: {min_val_bpb:.4f}")
+
+# Log to report
+if not dry_run:
+    from nanochat.report import get_report
+    get_report().log(section="Midtraining", data=[
+        user_config, # CLI args
+        { # stats about the training setup
+            "Number of iterations": step,
+            "DDP world size": ddp_world_size,
+        },
+        { # stats about training outcomes
+            "Minimum validation bpb": min_val_bpb,
+        }
+    ])
+
+# cleanup
+wandb_run.finish() # wandb run finish
+compute_cleanup()

scripts/tok_eval.py

@@ -0,0 +1,265 @@
+"""
+Evaluate compression ratio of the tokenizer.
+"""
+
+from nanochat.tokenizer import get_tokenizer, RustBPETokenizer
+from nanochat.dataset import parquets_iter_batched
+
+# Random text I got from a random website this morning
+news_text = r"""
+(Washington, D.C., July 9, 2025)- Yesterday, Mexico’s National Service of Agro-Alimentary Health, Safety, and Quality (SENASICA) reported a new case of New World Screwworm (NWS) in Ixhuatlan de Madero, Veracruz in Mexico, which is approximately 160 miles northward of the current sterile fly dispersal grid, on the eastern side of the country and 370 miles south of the U.S./Mexico border. This new northward detection comes approximately two months after northern detections were reported in Oaxaca and Veracruz, less than 700 miles away from the U.S. border, which triggered the closure of our ports to Mexican cattle, bison, and horses on May 11, 2025.
+
+While USDA announced a risk-based phased port re-opening strategy for cattle, bison, and equine from Mexico beginning as early as July 7, 2025, this newly reported NWS case raises significant concern about the previously reported information shared by Mexican officials and severely compromises the outlined port reopening schedule of five ports from July 7-September 15. Therefore, in order to protect American livestock and our nation’s food supply, Secretary Rollins has ordered the closure of livestock trade through southern ports of entry effective immediately.
+
+“The United States has promised to be vigilant — and after detecting this new NWS case, we are pausing the planned port reopening’s to further quarantine and target this deadly pest in Mexico. We must see additional progress combatting NWS in Veracruz and other nearby Mexican states in order to reopen livestock ports along the Southern border,” said U.S. Secretary of Agriculture Brooke L. Rollins. “Thanks to the aggressive monitoring by USDA staff in the U.S. and in Mexico, we have been able to take quick and decisive action to respond to the spread of this deadly pest.”
+""".strip()
+
+# Random Korean text (to test non-English compression)
+korean_text = r"""
+정직한 사실 위에, 공정한 시선을 더하다
+Herald Korea Times
+
+헤럴드코리아타임즈는 정치, 경제, 사회, 문화 등 한국 사회 전반의 주요 이슈를 심도 있게 다루는 종합 온라인 신문사입니다.
+
+우리는 단순히 뉴스를 전달하는 것이 아니라, 사실(Fact)에 기반한 양측의 시각을 균형 있게 조명하며, 독자 여러분이 스스로 판단할 수 있는 ‘정보의 균형’을 제공합니다.
+
+한국 언론의 오랜 문제로 지적되어 온 정치적 편향, 이념적 왜곡에서 벗어나
+오직 정직함과 공정함을 원칙으로 삼는 언론을 지향합니다.
+어느 한쪽의 주장만을 확대하거나 감추지 않고,
+**모든 쟁점에 대해 ‘무엇이 쟁점인지’, ‘누가 무엇을 주장하는지’, ‘사실은 무엇인지’**를 명확히 전달하는 데 집중합니다.
+""".strip()
+
+# Random piece of code
+code_text = r"""
+class BasicTokenizer(Tokenizer):
+
+    def __init__(self):
+        super().__init__()
+
+    def train(self, text, vocab_size, verbose=False):
+        assert vocab_size >= 256
+        num_merges = vocab_size - 256
+
+        # input text preprocessing
+        text_bytes = text.encode("utf-8") # raw bytes
+        ids = list(text_bytes) # list of integers in range 0..255
+
+        # iteratively merge the most common pairs to create new tokens
+        merges = {} # (int, int) -> int
+        vocab = {idx: bytes([idx]) for idx in range(256)} # int -> bytes
+        for i in range(num_merges):
+            # count up the number of times every consecutive pair appears
+            stats = get_stats(ids)
+            # find the pair with the highest count
+            pair = max(stats, key=stats.get)
+            # mint a new token: assign it the next available id
+            idx = 256 + i
+            # replace all occurrences of pair in ids with idx
+            ids = merge(ids, pair, idx)
+            # save the merge
+            merges[pair] = idx
+            vocab[idx] = vocab[pair[0]] + vocab[pair[1]]
+            # prints
+            if verbose:
+                print(f"merge {i+1}/{num_merges}: {pair} -> {idx} ({vocab[idx]}) had {stats[pair]} occurrences")
+""".strip()
+
+math_text = r"""
+\documentclass[12pt]{article}
+\usepackage{amsmath,amsthm,amssymb}
+\usepackage[margin=1in]{geometry}
+
+\newtheorem{theorem}{Theorem}
+\newtheorem*{remark}{Remark}
+
+\begin{document}
+
+\begin{center}
+{\Large A Cute Identity: The Sum of Cubes is a Square}
+\end{center}
+
+\begin{theorem}
+For every integer $n \ge 1$,
+\[
+\sum_{k=1}^{n} k^{3} \;=\; \left(\frac{n(n+1)}{2}\right)^{2}.
+\]
+\end{theorem}
+
+\begin{proof}[Proof 1 (Induction)]
+Let $S(n) = \sum_{k=1}^{n} k^3$. For $n=1$, $S(1)=1=(1\cdot 2/2)^2$, so the base case holds.
+
+Assume $S(n)=\big(\tfrac{n(n+1)}{2}\big)^2$ for some $n\ge 1$.
+Then
+\[
+S(n+1)
+= S(n) + (n+1)^3
+= \left(\frac{n(n+1)}{2}\right)^2 + (n+1)^3.
+\]
+Factor out $(n+1)^2$:
+\[
+S(n+1)
+= (n+1)^2\left( \frac{n^2}{4} + (n+1) \right)
+= (n+1)^2\left( \frac{n^2 + 4n + 4}{4} \right)
+= (n+1)^2\left( \frac{(n+2)^2}{4} \right).
+\]
+Thus
+\[
+S(n+1)=\left(\frac{(n+1)(n+2)}{2}\right)^2,
+\]
+which matches the claimed formula with $n$ replaced by $n+1$. By induction, the identity holds for all $n\ge 1$.
+\end{proof}
+
+\begin{proof}[Proof 2 (Algebraic telescoping)]
+Recall the binomial identity
+\[
+(k+1)^4 - k^4 = 4k^3 + 6k^2 + 4k + 1.
+\]
+Summing both sides from $k=0$ to $n$ telescopes:
+\[
+(n+1)^4 - 0^4
+= \sum_{k=0}^{n}\big(4k^3 + 6k^2 + 4k + 1\big)
+= 4\sum_{k=1}^{n}k^3 + 6\sum_{k=1}^{n}k^2 + 4\sum_{k=1}^{n}k + (n+1).
+\]
+Using the standard sums
+\[
+\sum_{k=1}^{n}k = \frac{n(n+1)}{2}
+\quad\text{and}\quad
+\sum_{k=1}^{n}k^2 = \frac{n(n+1)(2n+1)}{6},
+\]
+solve for $\sum_{k=1}^{n}k^3$ to get
+\[
+\sum_{k=1}^{n}k^3 = \left(\frac{n(n+1)}{2}\right)^2.
+\]
+\end{proof}
+
+\begin{remark}
+Geometrically, the identity says: ``adding up $1^3,2^3,\dots,n^3$ builds a perfect square’’—namely the square of the $n$th triangular number. This is why one sometimes calls it the \emph{sum-of-cubes is a square} phenomenon.
+\end{remark}
+
+\end{document}
+""".strip()
+
+science_text = r"""
+Photosynthesis is a photochemical energy transduction process in which light-harvesting pigment–protein complexes within the thylakoid membranes of oxygenic phototrophs absorb photons and initiate charge separation at the reaction center, driving the linear electron transport chain from water to NADP⁺ via photosystem II, the cytochrome b₆f complex, and photosystem I, concomitantly generating a trans-thylakoid proton motive force utilized by chloroplastic ATP synthase. The light-dependent reactions produce ATP and NADPH, which fuel the Calvin–Benson–Bassham cycle in the stroma, wherein ribulose-1,5-bisphosphate is carboxylated by ribulose-1,5-bisphosphate carboxylase/oxygenase (RuBisCO) to form 3-phosphoglycerate, subsequently reduced and regenerated through a series of enzymatic steps, enabling net assimilation of CO₂ into triose phosphates and ultimately carbohydrates. This process is tightly regulated by photoprotective mechanisms, redox feedback, and metabolite flux, representing a central biochemical pathway coupling solar energy capture to the biosphere’s primary productivity.
+""".strip()
+
+# The tokenizer was trained on data from earlier shards, so it has seen this data
+train_docs = next(parquets_iter_batched(split="train"))
+train_text = "\n".join(train_docs)
+val_docs = next(parquets_iter_batched(split="val"))
+val_text = "\n".join(val_docs)
+
+all_text = [
+    ("news", news_text),
+    ("korean", korean_text),
+    ("code", code_text),
+    ("math", math_text),
+    ("science", science_text),
+    ("fwe-train", train_text),
+]
+if val_text:
+    all_text.append(("fwe-val", val_text))
+
+# Try out current default compared to GPT-2 and GPT-4 tokenizers
+tokenizer_results = {}
+vocab_sizes = {}
+
+for tokenizer_name in ["gpt2", "gpt4", "ours"]:
+
+    if tokenizer_name == "gpt2":
+        tokenizer = RustBPETokenizer.from_pretrained("gpt2") # gpt-2 base model tokenizer
+    elif tokenizer_name == "gpt4":
+        tokenizer = RustBPETokenizer.from_pretrained("cl100k_base") # gpt-4 base model tokenizer
+    else:
+        tokenizer = get_tokenizer()
+
+    vocab_sizes[tokenizer_name] = tokenizer.get_vocab_size()
+    tokenizer_results[tokenizer_name] = {}
+
+    for name, text in all_text:
+        encoded = tokenizer.encode(text)
+        decoded = tokenizer.decode(encoded)
+        assert decoded == text
+
+        encoded_bytes = text.encode('utf-8')
+        ratio = len(encoded_bytes) / len(encoded)
+        tokenizer_results[tokenizer_name][name] = {
+            'bytes': len(encoded_bytes),
+            'tokens': len(encoded),
+            'ratio': ratio
+        }
+
+# ANSI color codes
+GREEN = '\033[92m'
+RED = '\033[91m'
+RESET = '\033[0m'
+
+# Print vocab sizes
+print(f"\nVocab sizes:")
+print(f"GPT-2: {vocab_sizes['gpt2']}")
+print(f"GPT-4: {vocab_sizes['gpt4']}")
+print(f"Ours: {vocab_sizes['ours']}")
+
+def print_comparison(baseline_name, baseline_results, ours_results, all_text):
+    """Print comparison table between baseline tokenizer and ours."""
+    print(f"\nComparison with {baseline_name}:")
+    print("=" * 95)
+    print(f"{'Text Type':<10} {'Bytes':<8} {baseline_name:<15} {'Ours':<15} {'Relative':<12} {'Better':<10}")
+    print(f"{'':10} {'':8} {'Tokens':<7} {'Ratio':<7} {'Tokens':<7} {'Ratio':<7} {'Diff %':<12}")
+    print("-" * 95)
+
+    for name, text in all_text:
+        baseline_data = baseline_results[name]
+        ours_data = ours_results[name]
+
+        # Calculate relative difference (positive means ours is better, negative means worse)
+        # Using tokens: fewer tokens is better, so we calculate (baseline_tokens - ours_tokens) / baseline_tokens
+        relative_diff = ((baseline_data['tokens'] - ours_data['tokens']) / baseline_data['tokens']) * 100
+
+        # Determine which has better compression (higher ratio = better)
+        if baseline_data['ratio'] > ours_data['ratio']:
+            baseline_color, ours_color = GREEN, RED
+            better = baseline_name
+            diff_color = RED
+        elif ours_data['ratio'] > baseline_data['ratio']:
+            baseline_color, ours_color = RED, GREEN
+            better = "Ours"
+            diff_color = GREEN
+        else:
+            baseline_color, ours_color = "", ""
+            better = "Tie"
+            diff_color = ""
+
+        print(f"{name:<10} {baseline_data['bytes']:<8} "
+              f"{baseline_color}{baseline_data['tokens']:<7}{RESET} "
+              f"{baseline_color}{baseline_data['ratio']:<7.2f}{RESET} "
+              f"{ours_color}{ours_data['tokens']:<7}{RESET} "
+              f"{ours_color}{ours_data['ratio']:<7.2f}{RESET} "
+              f"{diff_color}{relative_diff:+7.1f}%{RESET}     "
+              f"{better:<10}")
+
+# Print comparisons
+print_comparison("GPT-2", tokenizer_results['gpt2'], tokenizer_results['ours'], all_text)
+print_comparison("GPT-4", tokenizer_results['gpt4'], tokenizer_results['ours'], all_text)
+
+# Log to report
+from nanochat.report import get_report
+lines = []
+for baseline_name in ["GPT-2", "GPT-4"]:
+    baseline_key = baseline_name.lower().replace('-', '')
+    baseline_results = tokenizer_results[baseline_key]
+    ours_results = tokenizer_results['ours']
+    lines.append(f"### Comparison with {baseline_name}")
+    lines.append("")
+    lines.append("| Text Type | Bytes | " + baseline_name + " Tokens | " + baseline_name + " Ratio | Ours Tokens | Ours Ratio | Relative Diff % |")
+    lines.append("|-----------|-------|--------------|--------------|-------------|------------|-----------------|")
+    for name, text in all_text:
+        baseline_data = baseline_results[name]
+        ours_data = ours_results[name]
+        relative_diff = ((baseline_data['tokens'] - ours_data['tokens']) / baseline_data['tokens']) * 100
+        lines.append(f"| {name} | {baseline_data['bytes']} | {baseline_data['tokens']} | {baseline_data['ratio']:.2f} | {ours_data['tokens']} | {ours_data['ratio']:.2f} | {relative_diff:+.1f}% |")
+    lines.append("")
+report_markdown = "\n".join(lines)
+get_report().log(section="Tokenizer evaluation", data=[
+    report_markdown,
+])

scripts/tok_train.py

@@ -0,0 +1,106 @@
+"""
+Train a tokenizer using the HuggingFace Tokenizers library.
+In the style of GPT-4 tokenizer.
+"""
+import os
+import time
+import argparse
+import torch
+from nanochat.tokenizer import RustBPETokenizer
+from nanochat.common import get_base_dir
+from nanochat.dataset import parquets_iter_batched
+
+# -----------------------------------------------------------------------------
+# Parse command line arguments
+
+parser = argparse.ArgumentParser(description='Train a BPE tokenizer')
+parser.add_argument('--max_chars', type=int, default=10_000_000_000, help='Maximum characters to train on (default: 10B)')
+parser.add_argument('--doc_cap', type=int, default=10_000, help='Maximum characters per document (default: 10,000)')
+parser.add_argument('--vocab_size', type=int, default=65536, help='Vocabulary size (default: 65536 = 2^16)')
+args = parser.parse_args()
+print(f"max_chars: {args.max_chars:,}")
+print(f"doc_cap: {args.doc_cap:,}")
+print(f"vocab_size: {args.vocab_size:,}")
+
+# -----------------------------------------------------------------------------
+# Text iterator
+
+def text_iterator():
+    """
+    1) Flatten the batches into a single iterator
+    2) Crop every document to args.doc_cap characters
+    3) Break when we've seen args.max_chars characters
+    """
+    nchars = 0
+    for batch in parquets_iter_batched(split="train"):
+        for doc in batch:
+            doc_text = doc
+            if len(doc_text) > args.doc_cap:
+                doc_text = doc_text[:args.doc_cap]
+            nchars += len(doc_text)
+            yield doc_text
+            if nchars > args.max_chars:
+                return
+text_iter = text_iterator()
+
+# -----------------------------------------------------------------------------
+# Train the tokenizer
+t0 = time.time()
+tokenizer = RustBPETokenizer.train_from_iterator(text_iter, args.vocab_size)
+t1 = time.time()
+train_time = t1 - t0
+print(f"Training time: {train_time:.2f}s")
+
+# -----------------------------------------------------------------------------
+# Save the tokenizer to disk
+base_dir = get_base_dir()
+tokenizer_dir = os.path.join(base_dir, "tokenizer")
+tokenizer.save(tokenizer_dir)
+
+# -----------------------------------------------------------------------------
+# Quick inline sanity check
+test_text = """Hello world! This is a test.
+Numbers: 123, 4567, 89
+Contractions: I'm, you're, it's
+Special chars: @#$%^&*()
+Unicode: 你好世界 🌍"""
+encoded = tokenizer.encode(test_text)
+decoded = tokenizer.decode(encoded)
+assert decoded == test_text
+
+# -----------------------------------------------------------------------------
+# One more thing: we wish to cache a mapping from token id to number of bytes of that token
+# for efficient evaluation of bits per byte. Unlike the typical mean loss, this
+# allows us to report a loss that is invariant to the vocab size of the tokenizer.
+# The bits per byte on the validation set is then one of the primary metrics we care about.
+vocab_size = tokenizer.get_vocab_size()
+special_set = set(tokenizer.get_special_tokens())
+token_strings = [tokenizer.decode([token_id]) for token_id in range(vocab_size)]
+token_bytes = []
+for token_id in range(vocab_size):
+    token_str = token_strings[token_id] # the Python string representation of this token
+    if token_str in special_set:
+        token_bytes.append(0) # special characters are not counted
+    else:
+        id_bytes = len(token_str.encode("utf-8")) # number of bytes that make up this token
+        token_bytes.append(id_bytes)
+token_bytes = torch.tensor(token_bytes, dtype=torch.int32, device='cpu')
+token_bytes_path = os.path.join(tokenizer_dir, "token_bytes.pt")
+with open(token_bytes_path, "wb") as f:
+    torch.save(token_bytes, f)
+print(f"Saved token_bytes to {token_bytes_path}")
+
+# Log to report
+from nanochat.report import get_report
+token_bytes_nonzero = (token_bytes[token_bytes > 0]).to(dtype=torch.float32)
+get_report().log(section="Tokenizer training", data=[
+    vars(args), # argparse command line arguments
+    {"train_time": train_time},
+    {"num_special_tokens": len(special_set)},
+    {
+        "token_bytes_min": int(token_bytes_nonzero.min().item()),
+        "token_bytes_max": int(token_bytes_nonzero.max().item()),
+        "token_bytes_mean": token_bytes_nonzero.mean().item(),
+        "token_bytes_std": token_bytes_nonzero.std().item(),
+    }
+])

speedrun.sh

@@ -0,0 +1,131 @@
+#!/bin/bash
+
+# This script is the "Best ChatGPT clone that $100 can buy",
+# It is designed to run in ~4 hours on 8XH100 node at $3/GPU/hour.
+
+# 1) Example launch (simplest):
+# bash speedrun.sh
+# 2) Example launch in a screen session (because the run takes ~4 hours):
+# screen -L -Logfile speedrun.log -S speedrun bash speedrun.sh
+# 3) Example launch with wandb logging, but see below for setting up wandb first:
+# WANDB_RUN=speedrun screen -L -Logfile speedrun.log -S speedrun bash speedrun.sh
+
+# Default intermediate artifacts directory is in ~/.cache/nanochat
+export OMP_NUM_THREADS=1
+export NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
+mkdir -p $NANOCHAT_BASE_DIR
+
+# -----------------------------------------------------------------------------
+# Python venv setup with uv
+
+# install uv (if not already installed)
+command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
+# create a .venv local virtual environment (if it doesn't exist)
+[ -d ".venv" ] || uv venv
+# install the repo dependencies
+uv sync --extra gpu
+# activate venv so that `python` uses the project's venv instead of system python
+source .venv/bin/activate
+
+# -----------------------------------------------------------------------------
+# wandb setup
+# If you wish to use wandb for logging (it's nice!, recommended).
+# 1) Make sure to first log in to wandb, e.g. run:
+#    `wandb login`
+# 2) Set the WANDB_RUN environment variable when running this script, e.g.:
+#    `WANDB_RUN=d26 bash speedrun.sh`
+if [ -z "$WANDB_RUN" ]; then
+    # by default use "dummy" : it's handled as a special case, skips logging to wandb
+    WANDB_RUN=dummy
+fi
+
+# -----------------------------------------------------------------------------
+# During the course of the run, we will be writing markdown reports to the report/
+# directory in the base dir. This command clears it out and writes a header section
+# with a bunch of system info and a timestamp that marks the start of the run.
+python -m nanochat.report reset
+
+# -----------------------------------------------------------------------------
+# Tokenizer
+
+# Install Rust / Cargo
+curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+source "$HOME/.cargo/env"
+
+# Build the rustbpe Tokenizer
+uv run maturin develop --release --manifest-path rustbpe/Cargo.toml
+
+# Download the first ~2B characters of pretraining dataset
+# look at dev/repackage_data_reference.py for details on how this data was prepared
+# each data shard is ~250M chars
+# so we download 2e9 / 250e6 = 8 data shards at this point
+# each shard is ~100MB of text (compressed), so this is about ~800MB of data on disk
+python -m nanochat.dataset -n 8
+# Immediately also kick off downloading more shards in the background while tokenizer trains
+# See comment below for why 240 is the right number here
+python -m nanochat.dataset -n 240 &
+DATASET_DOWNLOAD_PID=$!
+# train the tokenizer with vocab size 2**16 = 65536 on ~2B characters of data
+python -m scripts.tok_train --max_chars=2000000000
+# evaluate the tokenizer (report compression ratio etc.)
+python -m scripts.tok_eval
+
+# -----------------------------------------------------------------------------
+# Base model (pretraining)
+
+# The d20 model is 561M parameters.
+# Chinchilla says #tokens = 20X #params, so we need 561e6 * 20 = 11.2B tokens.
+# Assume our tokenizer is 4.8 chars/token, this is 11.2B * 4.8 ~= 54B chars.
+# At 250M chars/shard, this is 54B / 250M ~= 216 shards needed for pretraining.
+# Round up to 240 for safety. At ~100MB/shard, this downloads ~24GB of data to disk.
+# (The total number of shards available in the entire dataset is 1822.)
+echo "Waiting for dataset download to complete..."
+wait $DATASET_DOWNLOAD_PID
+
+# Number of processes/GPUs to use
+NPROC_PER_NODE=8
+
+# pretrain the d20 model
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_train -- --depth=20 --run=$WANDB_RUN
+# evaluate the model on a larger chunk of train/val data and draw some samples
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_loss
+# evaluate the model on CORE tasks
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_eval
+
+# -----------------------------------------------------------------------------
+# Midtraining (teach the model conversation special tokens, tool use, multiple choice)
+
+# download 2.3MB of synthetic identity conversations to impart a personality to nanochat
+# see dev/gen_sft_data.py for details on how this data was prepared and to get a sense of how you can easily tune it
+curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl
+
+# run midtraining and eval the model
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.mid_train -- --run=$WANDB_RUN
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_eval -- -i mid
+
+# -----------------------------------------------------------------------------
+# Supervised Finetuning (domain adaptation to each sequence all by itself per row)
+
+# train sft and re-eval right away (should see a small bump)
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_sft -- --run=$WANDB_RUN
+torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_eval -- -i sft
+
+# chat with the model over CLI! Leave out the -p to chat interactively
+# python -m scripts.chat_cli -p "Why is the sky blue?"
+
+# even better, chat with your model over a pretty WebUI ChatGPT style
+# python -m scripts.chat_web
+
+# -----------------------------------------------------------------------------
+# Reinforcement Learning. Optional, and currently only on GSM8K
+# (optional)
+
+# run reinforcement learning
+# torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_rl -- --run=$WANDB_RUN
+# eval the RL model only on GSM8K
+# torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.chat_eval -- -i rl -a GSM8K
+
+# -----------------------------------------------------------------------------
+# Generate the full report by putting together all the sections
+# report.md is the output and will be copied to current directory for convenience
+python -m nanochat.report generate

tasks/arc.py

@@ -0,0 +1,48 @@
+"""
+The ARC dataset from Allen AI.
+https://huggingface.co/datasets/allenai/ai2_arc
+"""
+
+from datasets import load_dataset
+from tasks.common import Task, render_mc
+
+class ARC(Task):
+
+    def __init__(self, subset, split, **kwargs):
+        super().__init__(**kwargs)
+        assert subset in ["ARC-Easy", "ARC-Challenge"], "ARC subset must be ARC-Easy or ARC-Challenge"
+        assert split in ["train", "validation", "test"], "ARC split must be train|validation|test"
+        self.ds = load_dataset("allenai/ai2_arc", subset, split=split).shuffle(seed=42)
+
+    @property
+    def eval_type(self):
+        return 'categorical'
+
+    def num_examples(self):
+        return len(self.ds)
+
+    def get_example(self, index):
+        row = self.ds[index]
+        question = row["question"] # the question text
+        choices = row["choices"]["text"] # the text of each choice
+        answer_string = row["answerKey"] # e.g. "A", "B", "C", "D"
+        letters = row["choices"]["label"] # e.g. ["A", "B", "C", "D"]
+        assert answer_string in letters, f"ARC answer {answer_string} must be one of {letters}" # sanity check
+        # create and return the Conversation object
+        user_message = render_mc(question, letters, choices)
+        messages = [
+            {"role": "user", "content": user_message},
+            {"role": "assistant", "content": answer_string}
+        ]
+        conversation = {
+            "messages": messages,
+            "letters": letters, # useful during evaluation, so we can narrow and clamp the assistant prediction to one of the letters
+        }
+        return conversation
+
+    def evaluate(self, conversation, assistant_response):
+        # the assert here is not strictly speaking needed, but currently the way we eval, we expect this to be true
+        # I'm going to leave the assert here to prevent footguns, but possibly in the future can remove it.
+        assert assistant_response in conversation['letters'], f"ARC answer {assistant_response} is expected to be one of {conversation['letters']}"
+        assistant_message = conversation['messages'][-1]['content'] # e.g. "A"
+        return assistant_response == assistant_message

tasks/common.py

@@ -0,0 +1,147 @@
+"""
+Base class for all Tasks.
+A Task is basically a dataset of conversations, together with some
+metadata and often also evaluation criteria.
+Example tasks: MMLU, ARC-Easy, ARC-Challenge, GSM8K, HumanEval, SmolTalk.
+"""
+
+import random
+
+class Task:
+    """
+    Base class of a Task. Allows for lightweight slicing of the underlying dataset.
+    """
+
+    def __init__(self, start=0, stop=None, step=1):
+        # allows a lightweight logical view over a dataset
+        assert start >= 0, f"Start must be non-negative, got {start}"
+        assert stop is None or stop >= start, f"Stop should be greater than or equal to start, got {stop} and {start}"
+        assert step >= 1, f"Step must be strictly positive, got {step}"
+        self.start = start
+        self.stop = stop # could be None here
+        self.step = step
+
+    @property
+    def eval_type(self):
+        # one of 'generative' | 'categorical'
+        raise NotImplementedError
+
+    def num_examples(self):
+        raise NotImplementedError
+
+    def get_example(self, index):
+        raise NotImplementedError
+
+    def __len__(self):
+        start = self.start
+        stop = self.num_examples() if self.stop is None else self.stop
+        step = self.step
+        span = stop - start
+        num = (span + step - 1) // step # ceil_div(span, step)
+        assert num >= 0, f"Negative number of examples???: {num}" # prevent footguns
+        return num
+
+    def __getitem__(self, index: int):
+        assert isinstance(index, int), f"Index must be an integer, got {type(index)}"
+        physical_index = self.start + index * self.step
+        conversation = self.get_example(physical_index)
+        return conversation
+
+    def evaluate(self, problem, completion):
+        raise NotImplementedError
+
+
+class TaskMixture(Task):
+    """
+    For SFT Training it becomes useful to train on a tax mixture of datasets.
+    Fun trick: if you wish to oversample any task, just pass it in multiple times in the list.
+    """
+
+    def __init__(self, tasks, **kwargs):
+        super().__init__(**kwargs)
+        # tasks is a list of Task objects
+        self.tasks = tasks
+        self.lengths = [len(task) for task in self.tasks]
+        self.num_conversations = sum(self.lengths)
+        # Build list of all (task_idx, local_idx) pairs
+        self.index_map = []
+        for task_idx, task_length in enumerate(self.lengths):
+            for local_idx in range(task_length):
+                self.index_map.append((task_idx, local_idx))
+        # Deterministically shuffle to mix tasks throughout training
+        rng = random.Random(42)
+        rng.shuffle(self.index_map)
+        # Note: this is not the most elegant or best solution, but it's ok for now
+
+    def num_examples(self):
+        return self.num_conversations
+
+    def get_example(self, index):
+        """
+        Access conversations according to a deterministic shuffle of all examples.
+        This ensures tasks are mixed throughout training, regardless of dataset size.
+        """
+        assert 0 <= index < self.num_conversations, f"Index {index} out of range for mixture with {self.num_conversations} conversations"
+        task_idx, local_idx = self.index_map[index]
+        return self.tasks[task_idx][local_idx]
+
+
+class TaskSequence(Task):
+    """
+    For SFT Training sometimes we want to sequentially train on a list of tasks.
+    This is useful for cases that require a training curriculum.
+    """
+
+    def __init__(self, tasks, **kwargs):
+        super().__init__(**kwargs)
+        self.tasks = tasks
+        self.lengths = [len(task) for task in self.tasks]
+        self.num_conversations = sum(self.lengths)
+
+    def num_examples(self):
+        return self.num_conversations
+
+    def get_example(self, index):
+        assert 0 <= index < self.num_conversations, f"Index {index} out of range for sequence with {self.num_conversations} conversations"
+        for task_idx, task_length in enumerate(self.lengths):
+            if index < task_length:
+                return self.tasks[task_idx][index]
+            index -= task_length
+
+
+def render_mc(question, letters, choices):
+    """
+    The common multiple choice rendering format we will use.
+
+    Note two important design decisions:
+    1)
+    Bigger models don't care as much, but smaller models prefer to have
+    the letter *after* the choice, which results in better binding.
+    2)
+    There is no whitespace between the delimiter (=) and the letter.
+    This is actually critical because the tokenizer has different token ids
+    for " A" vs. "A". The assistant responses will be just the letter itself,
+    i.e. "A", so it is important that here in the prompt it is the exact same
+    token, i.e. "A" with no whitespace before it. Again, bigger models don't care
+    about this too much, but smaller models do care about some of these details.
+    """
+    query = f"Multiple Choice question: {question}\n"
+    query += "".join([f"- {choice}={letter}\n" for letter, choice in zip(letters, choices)])
+    query += "\nRespond only with the letter of the correct answer."
+    return query
+
+
+if __name__ == "__main__":
+    # very lightweight test of slicing
+    from tasks.mmlu import MMLU
+
+    ds = MMLU(subset="auxiliary_train", split="train")
+    print("Length of MMLU: ", len(ds))
+    ex = ds[5]
+    print("5th example: ", ex)
+
+    ds = MMLU(subset="auxiliary_train", split="train", start=5, stop=10)
+    print("Length of sliced MMLU[5:10]: ", len(ds))
+    print("0th example of sliced MMLU: ", ds[0])
+
+    print("They match: ", ex == ds[0])

tasks/customjson.py

@@ -0,0 +1,65 @@
+"""
+CustomJSON task for loading conversations from JSONL files.
+Each line in the JSONL file should be a JSON array of messages.
+"""
+
+import os
+import json
+from tasks.common import Task
+
+class CustomJSON(Task):
+    """
+    Load conversations from a JSONL file.
+    Each line should be a JSON array of message objects with 'role' and 'content' fields.
+    Example line: [{"role":"user","content":"Hi"},{"role":"assistant","content":"Hello"}]
+    """
+
+    def __init__(self, filepath, **kwargs):
+        super().__init__(**kwargs)
+        self.filepath = filepath
+        self.conversations = []
+
+        # Load all conversations from the JSONL file
+        if not os.path.exists(filepath):
+            # Helpful error message due to recent change. Will be removed in the future.
+            print("-" * 80)
+            print(f"Warning: File {filepath} does not exist")
+            print("HINT (Oct 21 2025)")
+            print("If you recently did a git pull and suddely see this, it might be due to the new addition of identity conversations")
+            print("See this discussion for more details: https://github.com/karpathy/nanochat/discussions/139")
+            print("Quick fix: simply run the following command to download the file and you're done:")
+            print(f"curl -L -o {filepath} https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl")
+            print("-" * 80)
+
+        else:
+            with open(filepath, 'r', encoding='utf-8') as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:  # skip empty lines
+                        continue
+                    messages = json.loads(line)
+                    # Validate the conversation structure
+                    assert isinstance(messages, list), f"Expected list of messages, got {type(messages)}"
+                    assert len(messages) >= 2, f"Conversation must have at least 2 messages, got {len(messages)}"
+                    # Validate message structure and alternating roles
+                    for i, message in enumerate(messages):
+                        assert "role" in message, f"Message {i} missing 'role' field"
+                        assert "content" in message, f"Message {i} missing 'content' field"
+                        expected_role = "user" if i % 2 == 0 else "assistant"
+                        assert message["role"] == expected_role, f"Message {i} has role {message['role']} but should be {expected_role}"
+                        assert isinstance(message["content"], str), f"Message {i} content must be a string"
+
+                    self.conversations.append(messages)
+
+        self.length = len(self.conversations)
+
+    def num_examples(self):
+        return self.length
+
+    def get_example(self, index):
+        messages = self.conversations[index]
+        conversation = {
+            "messages": messages,
+        }
+        return conversation
+

tasks/gsm8k.py

@@ -0,0 +1,117 @@
+"""
+GSM8K evaluation.
+https://huggingface.co/datasets/openai/gsm8k
+
+Example problem instance:
+
+Question:
+Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?
+Answer:
+Weng earns 12/60 = $<<12/60=0.2>>0.2 per minute.
+Working 50 minutes, she earned 0.2 x 50 = $<<0.2*50=10>>10.
+#### 10
+
+Notice that GSM8K uses tool calls inside << >> tags.
+"""
+
+import re
+from datasets import load_dataset
+from tasks.common import Task
+
+
+GSM_RE = re.compile(r"#### (\-?[0-9\.\,]+)")
+def extract_answer(completion):
+    """
+    Extract the numerical answer after #### marker.
+    Follows official code for normalization:
+    https://github.com/openai/grade-school-math/blob/3101c7d5072418e28b9008a6636bde82a006892c/grade_school_math/dataset.py#L28
+    """
+    match = GSM_RE.search(completion)
+    if match:
+        match_str = match.group(1).strip()
+        match_str = match_str.replace(",", "")
+        return match_str
+    return None
+
+
+class GSM8K(Task):
+
+    def __init__(self, subset, split, **kwargs):
+        super().__init__(**kwargs)
+        assert subset in ["main", "socratic"], "GSM8K subset must be main|socratic"
+        assert split in ["train", "test"], "GSM8K split must be train|test"
+        self.ds = load_dataset("openai/gsm8k", subset, split=split).shuffle(seed=42)
+
+    @property
+    def eval_type(self):
+        return 'generative'
+
+    def num_examples(self):
+        return len(self.ds)
+
+    def get_example(self, index):
+        """ Get a single problem from the dataset. """
+        row = self.ds[index]
+        question = row['question'] # string of the question prompt
+        answer = row['answer'] # string of the full solution and the answer after #### marker
+        # Create and return the Conversation object
+        # This is tricky because GSM8K uses tool calls, which we need to parse here.
+        assistant_message_parts = []
+        parts = re.split(r'(<<[^>]+>>)', answer)
+        for part in parts:
+            if part.startswith('<<') and part.endswith('>>'):
+                # This is a calculator tool call
+                inner = part[2:-2]  # Remove << >>
+                # Split on = to get expression and result
+                if '=' in inner:
+                    expr, result = inner.rsplit('=', 1)
+                else:
+                    expr, result = inner, ""
+                # Add the tool call as a part
+                assistant_message_parts.append({"type": "python", "text": expr})
+                # Add the result as a part
+                assistant_message_parts.append({"type": "python_output", "text": result})
+            else:
+                # Regular text in between tool calls
+                assistant_message_parts.append({"type": "text", "text": part})
+        # No put it all together
+        messages = [
+            {"role": "user", "content": question}, # note: simple string
+            {"role": "assistant", "content": assistant_message_parts}, # note: list of parts (as dicts)
+        ]
+        conversation = {
+            "messages": messages,
+        }
+        return conversation
+
+    def evaluate(self, conversation, assistant_response):
+        """
+        Given (conversation, completion), return evaluation outcome (0 = wrong, 1 = correct)
+        Note that:
+        - the conversation has both user AND assistant message (containing the ground truth answer)
+        - the assistant_response is usually the alternative assistant message achieved via sampling
+
+        TODO: Technically, assistant_response should be a Message (either a string or a list of parts)
+              We can handle this later possibly. For now just assume string.
+        """
+        assert isinstance(assistant_response, str), "Assuming simple string response for now"
+        # First extract the ground truth answer
+        assistant_message = conversation['messages'][-1]
+        assert assistant_message['role'] == "assistant", "Last message must be from the Assistant"
+        assert isinstance(assistant_message['content'], list), "This is expected to be a list of parts"
+        last_text_part = assistant_message['content'][-1]['text'] # this contains the final answer in GSM8K
+        # Extract both the ground truth answer and the predicted answer
+        ref_num = extract_answer(last_text_part)
+        pred_num = extract_answer(assistant_response)
+        # Compare and return the success as int
+        is_correct = int(pred_num == ref_num)
+        return is_correct
+
+    def reward(self, conversation, assistant_response):
+        """
+        Used during RL. To keep things simple, just re-use the evaluation above.
+        Later this could be made more complex (e.g. format matching etc.)
+        """
+        is_correct = self.evaluate(conversation, assistant_response)
+        is_correct_float = float(is_correct)
+        return is_correct_float