Repository: karpathy/nanochat Branch: master Commit: 5019accc5bc7 Files: 52 Total size: 606.5 KB Directory structure: gitextract_21_vu5fa/ ├── .claude/ │ └── skills/ │ └── read-arxiv-paper/ │ └── SKILL.md ├── .gitignore ├── .python-version ├── LICENSE ├── README.md ├── dev/ │ ├── LEADERBOARD.md │ ├── LOG.md │ ├── estimate_gpt3_core.ipynb │ ├── gen_synthetic_data.py │ ├── generate_logo.html │ ├── repackage_data_reference.py │ └── scaling_analysis.ipynb ├── nanochat/ │ ├── __init__.py │ ├── checkpoint_manager.py │ ├── common.py │ ├── core_eval.py │ ├── dataloader.py │ ├── dataset.py │ ├── engine.py │ ├── execution.py │ ├── flash_attention.py │ ├── fp8.py │ ├── gpt.py │ ├── loss_eval.py │ ├── optim.py │ ├── report.py │ ├── tokenizer.py │ └── ui.html ├── pyproject.toml ├── runs/ │ ├── miniseries.sh │ ├── runcpu.sh │ ├── scaling_laws.sh │ └── speedrun.sh ├── scripts/ │ ├── base_eval.py │ ├── base_train.py │ ├── chat_cli.py │ ├── chat_eval.py │ ├── chat_rl.py │ ├── chat_sft.py │ ├── chat_web.py │ ├── tok_eval.py │ └── tok_train.py ├── tasks/ │ ├── arc.py │ ├── common.py │ ├── customjson.py │ ├── gsm8k.py │ ├── humaneval.py │ ├── mmlu.py │ ├── smoltalk.py │ └── spellingbee.py └── tests/ ├── test_attention_fallback.py └── test_engine.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .claude/skills/read-arxiv-paper/SKILL.md ================================================ --- name: read-arxiv-paper description: Use this skill when when asked to read an arxiv paper given an arxiv URL --- You will be given a URL of an arxiv paper, for example: https://www.arxiv.org/abs/2601.07372 ### Part 1: Normalize the URL The goal is to fetch the TeX Source of the paper (not the PDF!), the URL always looks like this: https://www.arxiv.org/src/2601.07372 Notice the /src/ in the url. Once you have the URL: ### Part 2: Download the paper source Fetch the url to a local .tar.gz file. A good location is `~/.cache/nanochat/knowledge/{arxiv_id}.tar.gz`. (If the file already exists, there is no need to re-download it). ### Part 3: Unpack the file in that folder Unpack the contents into `~/.cache/nanochat/knowledge/{arxiv_id}` directory. ### Part 4: Locate the entrypoint Every latex source usually has an entrypoint, such as `main.tex` or something like that. ### Part 5: Read the paper Once you've found the entrypoint, Read the contents and then recurse through all other relevant source files to read the paper. #### Part 6: Report Once you've read the paper, produce a summary of the paper into a markdown file at `./knowledge/summary_{tag}.md`. Notice that 1) use the local knowledge directory here (it's easier for me to open and reference here), not in `~/.cache`, and 2) generate some reasonable `tag` like e.g. `conditional_memory` or whatever seems appropriate given the paper. Probably make sure that the tag doesn't exist yet so you're not overwriting files. As for the summary itself, remember that you're processing this paper within the context of the nanochat repository, so most often we we will be interested in how to apply the paper and its lessons to the nanochat project. Therefore, you should feel free to "remind yourself" of the related nanochat code by reading the relevant parts, and then explicitly make the connection of how this paper might relate to nanochat or what are things we might be inspired about or try. ================================================ FILE: .gitignore ================================================ .venv/ __pycache__/ *.pyc dev-ignore/ report.md eval_bundle/ # Secrets .env # Local setup CLAUDE.md wandb/ ================================================ FILE: .python-version ================================================ 3.10 ================================================ FILE: LICENSE ================================================ 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. ================================================ FILE: README.md ================================================ # nanochat ![nanochat logo](dev/nanochat.png) ![scaling laws](dev/scaling_laws_jan26.png) nanochat is the simplest experimental harness for training LLMs. It is designed to run on a single GPU node, the code is minimal/hackable, and it covers all major LLM stages including tokenization, pretraining, finetuning, evaluation, inference, and a chat UI. For example, you can train your own GPT-2 capability LLM (which cost ~$43,000 to train in 2019) for only $48 (~2 hours of 8XH100 GPU node) and then talk to it in a familiar ChatGPT-like web UI. On a spot instance, the total cost can be closer to ~$15. More generally, nanochat is configured out of the box to train an entire miniseries of compute-optimal models by setting one single complexity dial: `--depth`, the number of layers in the GPT transformer model (GPT-2 capability happens to be approximately depth 26). All other hyperparameters (the width of the transformer, number of heads, learning rate adjustments, training horizons, weight decays, ...) are calculated automatically in an optimal way. For questions about the repo, I recommend either using [DeepWiki](https://deepwiki.com/karpathy/nanochat) from Devin/Cognition to ask questions about the repo, or use the [Discussions tab](https://github.com/karpathy/nanochat/discussions), or come by the [#nanochat](https://discord.com/channels/1020383067459821711/1427295580895314031) channel on Discord. ## Time-to-GPT-2 Leaderboard Presently, the main focus of development is on tuning the pretraining stage, which takes the most amount of compute. Inspired by the modded-nanogpt repo and to incentivise progress and community collaboration, nanochat maintains a leaderboard for a "GPT-2 speedrun", which is the wall-clock time required to train a nanochat model to GPT-2 grade capability, as measured by the DCLM CORE score. The [runs/speedrun.sh](runs/speedrun.sh) script always reflects the reference way to train GPT-2 grade model and talk to it. The current leaderboard looks as follows: | # | time | val_bpb | CORE | Description | Date | Commit | Contributors | |---|-------------|---------|------|-------------|------|--------|--------------| | 0 | 168 hours | - | 0.2565 | Original OpenAI GPT-2 checkpoint | 2019 | - | OpenAI | | 1 | 3.04 | 0.74833 | 0.2585 | d24 baseline, slightly overtrained | Jan 29 2026 | 348fbb3 | @karpathy | | 2 | 2.91 | 0.74504 | 0.2578 | d26 slightly undertrained **+fp8** | Feb 2 2026 | a67eba3 | @karpathy | | 3 | 2.76 | 0.74645 | 0.2602 | bump total batch size to 1M tokens | Feb 5 2026 | 2c062aa | @karpathy | | 4 | 2.02 | 0.71854 | 0.2571 | change dataset to NVIDIA ClimbMix | Mar 4 2026 | 324e69c | @ddudek @karpathy | | 5 | 1.80 | 0.71808 | 0.2690 | autoresearch [round 1](https://x.com/karpathy/status/2031135152349524125) | Mar 9 2026 | 6ed7d1d | @karpathy | | 5 | 1.65 | 0.71800 | 0.2626 | autoresearch round 2 | Mar 14 2026 | a825e63 | @karpathy | The primary metric we care about is "time to GPT-2" - the wall clock time needed to outperform the GPT-2 (1.6B) CORE metric on an 8XH100 GPU node. The GPT-2 CORE score is 0.256525. In 2019, the training of GPT-2 cost approximately $43,000 so it is incredible that due to many advances over 7 years across the stack, we can now do so much faster and for well below $100 (e.g. at the current ~$3/GPU/hr, an 8XH100 node is ~$24/hr, so 2 hours is ~$48). See [dev/LEADERBOARD.md](dev/LEADERBOARD.md) for more docs on how to interpret and contribute to the leaderboard. ## Getting started ### Reproduce and talk to GPT-2 The most fun you can have is to train your own GPT-2 and talk to it. The entire pipeline to do so is contained in the single file [runs/speedrun.sh](runs/speedrun.sh), which is designed to be run on an 8XH100 GPU node. 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 runs/speedrun.sh ``` You may wish to do so in a screen session as this will take ~3 hours to run. Once it's done, you can talk to it 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 :). --- image --- A few more notes: - 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 personally exercised all of these code paths so there might be sharp edges. ## Research If you are a researcher and wish to help improve nanochat, two scripts of interest are [runs/scaling_laws.sh](runs/scaling_laws.sh) and [runs/miniseries.sh](runs/miniseries.sh). See [Jan 7 miniseries v1](https://github.com/karpathy/nanochat/discussions/420) for related documentation. For quick experimentation (~5 min pretraining runs) my favorite scale is to train a 12-layer model (GPT-1 sized), e.g. like this: ``` OMP_NUM_THREADS=1 torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- \ --depth=12 \ --run="d12" \ --model-tag="d12" \ --core-metric-every=999999 \ --sample-every=-1 \ --save-every=-1 \ ``` This uses wandb (run name "d12"), only runs the CORE metric on last step, and it doesn't sample and save intermediate checkpoints. I like to change something in the code, re-run a d12 (or a d16 etc) and see if it helped, in an iteration loop. To see if a run helps, I like to monitor the wandb plots for: 1. `val_bpb` (validation loss in vocab-size-invariant units of bits per byte) as a function of `step`, `total_training_time` and `total_training_flops`. 2. `core_metric` (the DCLM CORE socre) 3. VRAM utilization, `train/mfu` (Model FLOPS utilization), `train/tok_per_sec` (training throughput) See an example [here](https://github.com/karpathy/nanochat/pull/498#issuecomment-3850720044). The important thing to note is that nanochat is written and configured around one single dial of complexity - the depth of the transformer. This single integer automatically determines all other hyperparameters (the width of the transformer, number of heads, learning rate adjustments, training horizons, weight decays, ...) so that the trained model comes out compute optimal. The idea is that the user doesn't have to think about or set any of this, they are simply asking for a smaller or bigger model using `--depth`, and everything "just works". By sweeping out the depth, you achieve the nanochat miniseries of compute optimal models at various sizes. GPT-2 capability model (which is of most interest at the moment) happens to be somewhere around d24-d26 range with the current code. But any candidate changes to the repo have to be principled enough that they work for all settings of depth. ## Running on CPU / MPS The script [runs/runcpu.sh](runs/runcpu.sh) shows a very simple example of running on CPU or Apple Silicon. It dramatically shrinks the LLM that is being trained to make things fit into a reasonable time interval of a few ten minutes of training. You will not get strong results in this way. ## Precision / dtype nanochat does not use `torch.amp.autocast`. Instead, precision is managed explicitly through a single global `COMPUTE_DTYPE` (defined in `nanochat/common.py`). By default this is auto-detected based on your hardware: | Hardware | Default dtype | Why | |----------|--------------|-----| | CUDA SM 80+ (A100, H100, ...) | `bfloat16` | Native bf16 tensor cores | | CUDA SM < 80 (V100, T4, ...) | `float32` | No bf16; fp16 available via `NANOCHAT_DTYPE=float16` (uses GradScaler) | | CPU / MPS | `float32` | No reduced-precision tensor cores | You can override the default with the `NANOCHAT_DTYPE` environment variable: ```bash NANOCHAT_DTYPE=float32 python -m scripts.chat_cli -p "hello" # force fp32 NANOCHAT_DTYPE=bfloat16 torchrun --nproc_per_node=8 -m scripts.base_train # force bf16 ``` How it works: model weights are stored in fp32 (for optimizer precision), but our custom `Linear` layer casts them to `COMPUTE_DTYPE` during the forward pass. Embeddings are stored directly in `COMPUTE_DTYPE` to save memory. This gives us the same mixed-precision benefit as autocast but with full explicit control over what runs in which precision. Note: `float16` training automatically enables a `GradScaler` in `base_train.py` to prevent gradient underflow. SFT suppors this too but RL currently does not. Inference in fp16 works fine everywhere. ## Guides I've published a number of guides that might contain helpful information, most recent to least recent: - [Feb 1 2026: Beating GPT-2 for <<$100: the nanochat journey](https://github.com/karpathy/nanochat/discussions/481) - [Jan 7 miniseries v1](https://github.com/karpathy/nanochat/discussions/420) documents the first nanochat miniseries of models. - 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). - 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 the SFT stage. - [Oct 13 2025: original nanochat post](https://github.com/karpathy/nanochat/discussions/1) introducing nanochat, though now it contains some deprecated information and the model is a lot older (with worse results) than current master. ## 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 ├── nanochat │ ├── __init__.py # empty │ ├── checkpoint_manager.py # Save/Load model checkpoints │ ├── common.py # Misc small utilities, quality of life │ ├── 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) │ ├── optim.py # AdamW + Muon optimizer, 1GPU and distributed │ ├── 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 ├── runs │ ├── miniseries.sh # Miniseries training script │ ├── runcpu.sh # Small example of how to run on CPU/MPS │ ├── scaling_laws.sh # Scaling laws experiments │ └── speedrun.sh # Train the ~$100 nanochat d20 ├── scripts │ ├── base_eval.py # Base model: CORE score, bits per byte, samples │ ├── base_train.py # Base model: train │ ├── chat_cli.py # Chat model: talk to over CLI │ ├── chat_eval.py # Chat model: eval tasks │ ├── chat_rl.py # Chat model: reinforcement learning │ ├── chat_sft.py # Chat model: train SFT │ ├── chat_web.py # Chat model: talk to over WebUI │ ├── tok_eval.py # Tokenizer: evaluate compression rate │ └── tok_train.py # Tokenizer: train it ├── 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_engine.py └── uv.lock ``` ## Contributing The goal of nanochat 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 are 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 ChatGPT model you can talk to. Currently, the most interesting part personally is speeding up the latency to GPT-2 (i.e. getting a CORE score above 0.256525). Currently this takes ~3 hours, but by improving the pretraining stage we can improve this further. Current AI 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. - Thank you to the repo czar Sofie [@svlandeg](https://github.com/svlandeg) for help with managing issues, pull requests and discussions of nanochat. ## 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 ================================================ FILE: dev/LEADERBOARD.md ================================================ # Leaderboard Docs on participating in the "Time-to-GPT-2" leaderboard of nanochat. The primary metric we care about is "time to GPT-2" - the wall clock time needed to outperform the GPT-2 (1.6B) CORE metric on an 8XH100 GPU node. Originally in 2019, GPT-2 was trained by OpenAI on 32 TPU v3 chips for 168 hours (7 days), with $8/hour/TPUv3 back then, for a total cost of approx. $43K. It achieves 0.256525 CORE score, which is an ensemble metric introduced in the DCLM paper over 22 evaluations like ARC/MMLU/etc. ## How to The script [runs/speedrun.sh](runs/speedrun.sh) always implements the current state of the art on the leaderboard. In practice, I tune the base_train command a little bit. For example, once all the setup is configured and a tokenizer is trained, I like to do something like: ``` OMP_NUM_THREADS=1 torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- \ --depth=26 \ --run="d26-feb2-fp8-ratio8.25" \ --model-tag="d26_feb2_fp8_ratio8.25" \ --device-batch-size=16 \ --sample-every=-1 \ --save-every=-1 \ --core-metric-max-per-task=-1 \ --core-metric-every=999999 \ --target-param-data-ratio=8.25 \ --fp8 ``` Note that: - `depth` controls the size of the Transformer - `run` is the wandb name - `model-tag` is the location of the checkpoints on disk - `device-batch-size` in the ideal world, you want this to be 32 because with sequence length of 2048 (the default) and 8 GPUs we get `32 X 2048 X 8 = 524,288`, which is the total desired batch size determined to work fairly well around this scale. However, for bigger (e.g. d26), 32 is too much and OOMs, so we decrease it by 2 to 16. The `base_train.py` script automatically compensates for this by calculating that it has to use gradient accumulation of 2 to meet the desired total batch size. Therefore, it will do forward+backward twice and then a single step. Long story short, the ideal value is 32. If that doesn't fit, you decrease it, e.g. 16, 8, etc., keeping it powers of two so that the gradient accumulation math works out neatly. - `sample-every = -1` turns off periodic sampling - `core-metric-max-per-task=-1` means we run the entire CORE eval - `core-metric-every=999999` a bit of a hacky way to make the CORE eval only happen a single time at the very end of the run - `target-param-data-ratio=8.25` controls the training horizon, which is determined in the script by taking the number of non-embedding model parameters and simply multiplying by this number. The current optimal Tokens:Params ratio can be seen in the defaults of the `base_train.py` script (it is 10.5). 10.5 would produce the *compute optimal* model given the currently measured scaling laws. However, GPT-2 capability is currently somewhere in between a d24 and d26. So to reach it exactly, we want to either overtrain d24 or undertrain d26. In this particular example, I am choosing to slightly undertrain a d26. Note that odd depths (e.g. d25) are not super recommended to use because the math around the transformer sizing and its head dimensions doesn't come out neatly. - `--fp8` turns on fp8 training. If your GPU does not support fp8, you can leave this out and the code will simply train in bf16. bf16 is higher precision than fp8, so you can actually expect that you might be able to do fewer steps (lower the `target-param-data-ratio`) to achieve the same capability. Once you kick off the run, you wait ~3 hours and then at the end you'll see something like: ``` wandb: Run summary: wandb: core_metric 0.25851 wandb: step 16704 wandb: total_training_flops 4.330784131228946e+19 wandb: total_training_time 10949.46713 ``` Your CORE metric must be greater than GPT-2 0.256525. Then you report the `total_training_time`, (e.g. 10949) which is the time of the training iterations alone, excluding all the evaluations and logging, in seconds. So here for example it is roughly 10949/60/60 ~= 3.04 hours. You should also note and report the validation bpb of your run because the CORE metric can be a little bit noisy. If you outperform GPT-2 and the time is less than current SOTA in the Leaderboard, you get to make a PR. In addition to raw gains, there are some qualitative and aesthetic considerations that go into whether your improvement is merged. For example, if it is gnarly or it significantly bloats the code, or it seems too esoteric, then we will weigh those things against the improvement demonstrated. Additionally, nanochat cares not only about targeting a single model, but an entire miniseries of models. So your change must be principled enough that it can easily generalize to other model depths, so that we can sweep out a miniseries. After you create the commit, to get the current short git commit hash: ``` git log -1 --format="%h" ``` ## Run 1 Achieved Jan 29 2026 on commit `348fbb3`. The launch command was ``` OMP_NUM_THREADS=1 torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- \ --depth=24 \ --run=d24-jan29 \ --model-tag=d24_jan29 \ --device-batch-size=16 \ --sample-every=-1 \ --save-every=-1 \ --core-metric-max-per-task=-1 \ --core-metric-every=3000 \ --target-param-data-ratio=12 ``` The result was: ``` wandb: Run summary: wandb: core_metric 0.25851 wandb: step 16704 wandb: total_training_flops 4.330784131228946e+19 wandb: total_training_time 10949.46713 ``` The validation bpb was 0.74833. Detailed writeup: [Beating GPT-2 for <<$100: the nanochat journey](https://github.com/karpathy/nanochat/discussions/481) ## Run 2 Achieved Feb 2 2026 on commit `a67eba3`. The launch command was ``` OMP_NUM_THREADS=1 torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- \ --depth=26 \ --run="d26-feb2-fp8-ratio8.5" \ --model-tag="d26_feb2_fp8_ratio8.5" \ --device-batch-size=16 \ --sample-every=-1 \ --save-every=-1 \ --core-metric-max-per-task=-1 \ --core-metric-every=999999 \ --target-param-data-ratio=8.5 \ --fp8 ``` The result was: ``` core_metric 0.2578 step 14889 total_training_time 10493 Minimum validation bpb: 0.745036 ``` The big change in this run is `--fp8`, which causes all Linear layers (other than the gates) to be switched to fp8 training using `torchao` with tensorwise fp8 scaling. Each step is of slightly lower quality, but we are taking them a lot faster, coming out net ahead. Anyone who does not have fp8 (e.g. using a GPU without it) can simply leave out the `--fp8` flag to train in bfloat16. This will work just fine but it will produce a slightly stronger model than GPT-2 because of the fp8 -> bf16 precision upgrade. It's possible that one can further tune which layers to include in the fp8 conversion and that e.g. some of the smaller matmuls should be just kept in bf16 etc. Previous record was 3.04 hours, so 2.91 hours is `(3.04 - 2.91)/3.04*100` ~= 4.3% speed improvement. ## Run 3 Achieved Feb 5 2026 on commit `2c062aa`. Launch command: ``` OMP_NUM_THREADS=1 torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- \ --depth=26 \ --run="d26_feb4_double_batch_ratio8.25" \ --model-tag="d26_feb4_double_batch_ratio8.25" \ --device-batch-size=16 \ --total-batch-size=1048576 \ --sample-every=-1 \ --save-every=-1 \ --core-metric-max-per-task=-1 \ --core-metric-every=999999 \ --target-param-data-ratio=8.25 \ --fp8 ``` Result: ``` core_metric 0.26024 step 7226 total_training_time 9922 Minimum validation bpb: 0.74645 ``` The big change here is that the batch size was doubled from 0.5M to 1M, which works better for a d26 model and allowed me to decrease the number of optimization steps a bit via `--target-param-data-ratio` from 8.5 to 8.25. The TLDR is that the original batch size of 0.5M was tuned for d12, but bigger models (e.g. d26) prefer larger total batch size. I determined in experiments that d26 prefers 1M. Then I implemented and merged a principled way to calculate the optimal batch size given depth so that all nanochat models of all depths benefit. See [dev/LOG.md](dev/LOG.md) entry "2026-02-05: Auto Batch Size Scaling" for more detail. ## Run 4 Achived Mar 3 2026 on commit `324e69c`. The big change is the switch from HuggingFace FineWeb-EDU to NVIDIA ClimbMix dataset. `@karpathy` has tried to swap the dataset many times, each time with a negative result (FineWeb, DCLM, Olmo), but ClimbMix produced clear and immediate gains. Credit to `@ddudek` for originally discovering ClimbMix for nanochat and reporting the improvements, which kicked off the followup investigation. To reproduce, use the commit above, download at least 150 data shards, train the tokenizer: ``` python -m nanochat.dataset -n 150 python -m scripts.tok_train ``` Then kick off the run in the typical way, using a slightly lower than compute optimal ratio of 9.5 (vs compute optimal 10.5), meaning the d24 is slightly undertrained. ``` OMP_NUM_THREADS=1 torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- \ --depth=24 \ --run="d24-climbmix" \ --model-tag="d24-climbmix" \ --sample-every=-1 \ --save-every=-1 \ --core-metric-max-per-task=-1 \ --core-metric-every=999999 \ --target-param-data-ratio=9.5 \ --device-batch-size=16 \ --fp8 ``` I ran this command 7 individual times. Because our training is mildly non-deterministic, we get a spread of CORE scores, e.g.: ``` 0.25373 0.2584 0.25489 0.2568 0.25732 0.26765 0.25119 ``` Mean is 0.25714 (higher than the GPT-2 threshold needed), max-min is 0.01646. Something to investigate in the future is that even slightly better results can be obtained by randomly shuffling the the data shards (i.e. just going in a different order). This is unexpected because the documents were completely fully shuffled during data construction, so one would expect a relatively uniform data distribution. Indeed, the current default order is unfortunately among the worse ("unlucky") ones you can obtain with different shuffle seeds, but it suffices to beat GPT-2 for now so I am merging. TODO investing a bit more later. NOTE: The `val_bpb` is as of this run *NOT* comparable due to the data distribution change to the previous 3 runs. This run happens to be at `0.71854` validation bpb. If the dataset is not changed, the `val_bpb` number is a great, smooth metric to track relative performance w.r.t. and has less noise than CORE. ## Run 5 Achieved Mar 9, 2026 on commit `6ed7d1d`. Exactly the same launch command as Run 4 except `--target-param-data-ratio=8.7`. I ran 5 identical runs, the average CORE was 0.2690, which is quite a bit above the needed threshold of 0.2565. But the reason I didn't decrease the ratio further (i.e. train shorter) is that while the CORE "safety gap" is large, the val_loss safety gap is smaller - 0.71808, which we want to be below the Run 4 val loss of 0.71854. It's likely that we could have reduced the ratio even lower, possibly to 8.6, but it's not worth splitting hairs at this point. This commit is special because all of the improvements that went into [this commit](https://github.com/karpathy/nanochat/commit/6ed7d1d82cee16c2e26f45d559ad3338447a6c1b) came from fully autonomous "research" done by a private version of [autoresearch](https://github.com/karpathy/autoresearch) run on a d12 model. I wrote more about this in [this tweet](https://x.com/karpathy/status/2031135152349524125). The changes easily translated from d12 to d24, hence new leaderboard record, taking us from 2.02 hours "time to GPT-2" to 1.80 hours. ## Run 6 Achieved Mar 14, 2026 on commit `a825e63`. Exactly the same launch command as Run 4 except `--target-param-data-ratio=8`. Improvements in the architecture are allowing us to train shorter and shorter time. Instead of an undertrained d24 I attempted to train an overtrained d22 but it was worse. This set of changes came from autoresearch round 2, where I asked it to reference the modded-nanogpt repo for inspiration. So the exploration tried out a number of ideas and in particular found a way to incorporate the backout and smear in such a way that they are helpful (I had previously tried them manually a long time ago and they caused regressions). The smear idea in particular is a little bit heavier and bloaty because it is essentially an "early fusion" of context across tokens, producing a kind of a bigram input into the network and allowing it to focus on higher ngrams earlier. But for this reason the code gets a bit more complex and required some changes to inference. I verified with a unit test that the Engine inference is correct compared to the naive inference of `GPT.generate()`. The average of 5 runs was CORE 0.262634 and each of them lasted 1.65 hours (99 minutes). ================================================ FILE: dev/LOG.md ================================================ # Experiment Log A running summary documenting some experiments and findings. Started ~Jan 7 2026. --- ## 2026-03-04: Remove autocast, explicit dtype management, fp16 GradScaler Replaced `torch.amp.autocast` throughout the codebase with explicit dtype management via a single `COMPUTE_DTYPE` global. Also added fp16 training support with GradScaler. ### Motivation autocast is "magic we don't control" — it silently decides which ops run in which precision via internal allowlists. For this codebase, autocast was doing very little: the only thing it actually cast was `nn.Linear` weights from fp32 to bf16 for matmuls. `F.rms_norm`, `F.cross_entropy`, and Flash Attention all handle their own dtypes already. By making precision explicit, we gain fine-grained control (e.g. can experiment with fp32 norms) and eliminate an unnecessary layer of abstraction. ### What changed **Core mechanism** (`nanochat/common.py`, `nanochat/gpt.py`): - `COMPUTE_DTYPE` auto-detected from hardware: SM 80+ → bf16, pre-Ampere → fp32, CPU/MPS → fp32. Override via `NANOCHAT_DTYPE` env var. - Custom `Linear(nn.Linear)` class that casts weights to match input dtype in forward: `F.linear(x, self.weight.to(dtype=x.dtype))`. This is the single mechanism that replaces autocast. - Embeddings cast to `COMPUTE_DTYPE` at init (saves memory). Exception: fp16 keeps embeddings fp32 because GradScaler cannot unscale fp16 gradients. - Embedding output explicitly cast to `COMPUTE_DTYPE` in `GPT.forward()` (no-op for bf16, active for fp16 path). - RoPE cos/sin cache uses `COMPUTE_DTYPE` instead of hardcoded bf16. **Autocast removal** (11 files): - Deleted `--dtype` CLI flag, `ptdtype` variables, `autocast_ctx` definitions, and all `with autocast_ctx:` blocks from: `base_train.py`, `chat_sft.py`, `chat_rl.py`, `chat_cli.py`, `chat_eval.py`, `chat_web.py`, `base_eval.py`, `engine.py`, `bench_train_toks.py`, `test_e2e_pipeline.py`. **fp16 + GradScaler** (`base_train.py`, `chat_sft.py`): - `scaler = torch.amp.GradScaler() if COMPUTE_DTYPE == torch.float16 else None` - Backward: `scaler.scale(loss).backward()` vs plain `loss.backward()` - After accumulation: `scaler.unscale_(optimizer)` → distributed inf-sync via `scaler._found_inf_per_device(optimizer)` all-reduced with `ReduceOp.MAX` → `scaler.step(optimizer)` → `scaler.update()` - Zero overhead for bf16/fp32 paths (scaler is None, no branching inside kernels). **FP8 fix** (`nanochat/fp8.py`, `base_train.py`): - `Float8Linear.forward` explicitly casts input to `COMPUTE_DTYPE` (previously relied on autocast). - `disable_fp8` context manager now creates our custom `Linear` (not vanilla `nn.Linear`) when swapping out Float8Linear during eval. **Flash Attention** (`flash_attention.py`): - FA3 Hopper kernels don't support fp16 or fp32, so `USE_FA3` (module-level constant, resolved once at import) returns False, falling back to SDPA. --- ## 2026-03-04: Dataset upgrade: FineWeb-EDU 100B → ClimbMix 400B Switched the pretraining dataset from FineWeb-EDU 100B to ClimbMix 400B. This is by far the single biggest improvement to nanochat's GPT-2 speedrun time, bringing it down from **2 hours 46 minutes to 2 hours 1 minute** — a 27% reduction. ### What is ClimbMix? ClimbMix 400B is a curated 400B-token pretraining mixture hosted at `karpathy/climbmix-400b-shuffle` on HuggingFace. It comes form [NVIDIA](https://huggingface.co/datasets/nvidia/Nemotron-ClimbMix). It is a blend of high-quality web text, code, math, and other sources, designed to be a better general-purpose pretraining dataset than FineWeb-EDU alone. ### What changed - **Dataset**: `karpathy/fineweb-edu-100b-shuffle` → `karpathy/climbmix-400b-shuffle` (up to 6543 shards available vs the previous 1823 data shards, allowing for longer training in the future) - **Data directory**: `base_data/` → `base_data_climbmix/` (clean separation from legacy data) - **Model depth**: d26 → d24. ClimbMix trains more efficiently, so a smaller model reaches GPT-2 capability - **Shard count**: Only approx 150 data shards (~7B tokens) are now needed for GPT-2 capability - **Eval tokens**: doubled from 40 to 80 batches for more stable validation loss estimates - **Legacy fallback**: added a migration warning in `list_parquet_files()` that detects the old `base_data/` directory and falls back gracefully, so existing users see clear upgrade instructions on `git pull` ### Context This is the sixth attempt at beating FineWeb-EDU on CORE score — the previous five all failed (see entries on 2026-02-17, 2026-02-10, 2026-01-12 below). ClimbMix is the first dataset to convincingly surpass it, and the margin is large enough to also shrink the model from d26 to d24. --- ## 2026-03-02: SoftCap tuning Quick experiment to tune logit softcap on d24 scale. Tried 5..30. 5 was terrible, the rest of them were all about equal with the exception of 20, which was the best. Minor but solid improvement: val loss improved by ~1e-3 (0.716 -> 0.715). Setting as default. ## 2026-02-19: Mixture of Experts (negative) Implemented a DeepSeekV3-style Mixture of Experts layer as a drop-in replacement for the dense MLP. The MoE branch works and improves per-step validation loss, but is not a net improvement on wall clock time due to MoE overhead (at least for our scale of interest of approx GPT-2 capability). ### Implementation Follows DeepSeekV3 and using torchtitan as reference: - **8 routed experts, top-2 routing** with sigmoid gating (not softmax) - **1 shared expert** (dense MLP processing all tokens, following DeepSeekV3) - **Auxiliary-loss-free load balancing** (DeepSeekV3's expert bias nudging) - **Iso-FLOP sizing**: `expert_hidden_dim = round(4 * dim / (top_k + num_shared) / 128) * 128`, so active FLOPs per token match the dense MLP - **`torch._grouped_mm`** for dispatching tokens to experts in a single kernel (instead of a Python for-loop) - **3D expert weight tensors** `(num_experts, hidden, dim)` — Muon's Polar Express operates on the last two dims, so each expert is independently orthogonalized - **Active parameter counting** for scaling laws (only `top_k + shared` experts, not all 8) ### What was easy - The core MoE forward pass: router, sort tokens by expert, grouped matmul, scatter back. Conceptually clean. - Shared expert: just an `nn.Linear` MLP that runs on all tokens alongside the routed path. - 3D expert params + Muon: only required fixing `second_momentum_buffer` shape to preserve leading dims. - Load balancing: DeepSeekV3's bias nudging is simple and effective (~10 lines). ### What was hard / ugly - **`torch._grouped_mm` quirks**: requires bf16 (not fp32), column-major right operand, int32 cumulative offsets. The API is undocumented and only discoverable by trial and error. - **Token count padding**: torchtitan pads each expert's token count to alignment multiples (8 for bf16) for better grouped_mm throughput. We implemented this with both a pure PyTorch approach and a copy of torchtitan's Triton kernel. Both compiled cleanly (0 graph breaks), but with ~65K tokens across 8 experts, each expert already gets ~8K tokens which is well-aligned. The padding overhead (gather/scatter) actually regressed MFU from 35% to 33%. Reverted. - **FP8 + MoE**: `torch._grouped_mm` does NOT support FP8. There's a separate `torch._scaled_grouped_mm` API that requires per-row scaling (not per-tensor like our `Float8Linear`). The backward pass for weight gradients needs per-group column-wise scaling, which torchao implements with custom Triton kernels. We investigated thoroughly (see `dev/moe_fp8.md`) but did not implement — would require either depending on `torchao.prototype` (unstable) or writing ~200 lines of custom autograd + quantization code. Partial FP8 support exists: the shared expert's `nn.Linear` layers do get converted, but the routed experts (3D `nn.Parameter`) stay in bf16. ### Results - d18: MFU dropped from ~46% to ~35% (the grouped_mm dispatch + token sorting overhead is significant) - Per-step improvement in validation loss does not compensate for the throughput hit - Net negative on wall clock time ### What remains (if revisited) - **FP8 for routed experts**: Use `torch._scaled_grouped_mm` with a custom `_Float8GroupedMatmul` autograd function, with bf16 fallback for weight gradient (avoiding the per-group column-wise Triton kernels). What's really needed is a fused "FlashMoE" kernel that handles routing + expert dispatch + matmul in one shot (like FlashAttention did for attention), with all the needed features. This doesn't exist yet. Rawdogging MoE with current PyTorch primitives is painful — lots of sorting, gathering, scattering, and layout wrangling around the actual compute. ### Verdict MoE is not worth the trouble for nanochat right now. The code bloat is substantial (moe.py, router, shared expert, load balancing, optimizer fixes, FP8 gaps, active param counting) and the performance is worse wall-clock at our scale of interest. The fundamental issue is that the grouped_mm dispatch overhead eats the FLOP savings from sparsity, at least at our model scales and sequence lengths. --- ## 2026-02-17: Pretraining Data: FineWeb (negative) Tried vanilla fineweb instead of fineweb-edu dataset. Significantly, shockingly worse results: - d26 (GPT-2): CORE 0.2602 → 0.2241 This is the fifth failed attempt to beat pure FineWeb-EDU on CORE score. --- ## 2026-02-17: Pretraining Data Mixture Experiment (negative) Tried [hynky/finepdfs_50BT-dclm_30BT-fineweb_edu_20BT](https://huggingface.co/datasets/hynky/finepdfs_50BT-dclm_30BT-fineweb_edu_20BT), a mixture of FinePDFs, DCLM, and FineWeb-EDU. Slightly worse on both model sizes tested: - d26 (GPT-2): CORE 0.2602 → 0.2549 - d18: CORE 0.199 → 0.192 This is the fourth failed attempt to beat pure FineWeb-EDU on CORE score. --- ## 2026-02-16: SFT Script Upgrades Brought `chat_sft.py` up to parity with `base_train.py` and tuned settings based on SFT sweeps. Tuning: - **Optimizer warm-start** (`--load-optimizer=1`, default on): loads pretrained momentum buffers via new `load_optimizer_state()` in `checkpoint_manager.py`. LRs are reset to fresh SFT values after load. Loading the optimizer works slightly better but not by too much. - **LR schedule**: replaced "constant 80%, linear to 0" with warmup/constant/warmdown matching `base_train.py` (`--warmup-ratio`, `--warmdown-ratio`, `--init-lr-frac`, `--final-lr-frac`). Similar to pretraining, warmdown ratio of 0.5 worked the best. `--init-lr-frac` changed from 1.0 slightly lower to 0.8. - **LR tuning**: attempted to tune all the individual LRs (e.g. does SFT prefer lower LR for embeddings? etc.) but all of this produced negative results. - **Data mixture**: MMLU epochs 1→3, GSM8K epochs 2→4 (confirmed best from sweeps). Epoch counts now configurable via `--mmlu-epochs` / `--gsm8k-epochs`. Might remove these in the future though. Quality of life, footguns, minor fixes: - **Hyperparameter inheritance**: SFT now inherits batch sizes and LRs from the pretrained checkpoint metadata by default (CLI overrides still work). Also saved `total_batch_size` to `base_train.py` checkpoint metadata. - **GC management**: disabled Python GC after step 1 to avoid ~500ms pauses (manual collect every 5000 steps), same as base pretraining. - **ChatCORE eval**: periodic eval during SFT (`--chatcore-every=200`) across all 6 tasks, logged to wandb. - **MFU**: uses `get_peak_flops()` for actual GPU instead of hardcoded H100 value. - Removed `--dry-run` and `--dtype` flags. All ranks now participate in checkpoint save. --- ## 2026-02-05: Auto Batch Size Scaling ### Background So far, the `--total-batch-size` was hardcoded to be `2**19 = 524,288` ~= 0.5M tokens. This was the optimal setting for d12, but when I tried to re-tune it for d26 (GPT-2), I noticed that the optimal was closer to `2**20 = 1,048,576` ~= 1M tokens. This is to be expected - larger models prefer a higher optimal total batch size. However, we have to make sure that all settings of `--depth` get their own optimal batch size calculated in some principled way. Here, I referenced the "Power Lines" paper from Cerebras ([arXiv:2505.13738](https://arxiv.org/abs/2505.13738)) for a lot of related experimentation. In particular, they found that **Bopt ∝ D^0.383** (where D is the number of training tokens, not the number of parameters!). So the idea is to tune the optimal batch size on d12, and then extrapolate it with this power law to bigger models. The 0.383 exponent means batch size grows slowly: 10× more tokens only justifies ~2.4× bigger batch. For nanochat's compute-optimal training (D ∝ N via `--target-param-data-ratio`), this means deeper models naturally want larger batches. ### Implementation Added `--total-batch-size=-1` (now the default) to auto-compute optimal batch: ```python get_scaling_params = lambda m: m.num_scaling_params()['transformer_matrices'] + m.num_scaling_params()['lm_head'] if args.total_batch_size == -1: D_REF = args.target_param_data_ratio * get_scaling_params(build_model_meta(12)) B_REF = 2**19 args.total_batch_size = 2 ** round(math.log2(B_REF * (target_tokens / D_REF) ** 0.383)) ``` Reference point: d=12 model with B=2^19 (empirically validated). The reference is computed dynamically so that if the architecture changes (e.g., different `--aspect-ratio`), the math automatically adjusts. However, if the model actually does change too much, one would also want to re-tune the optimal batch size for d=12. ### Results With this formula, we currently get: | Depth | Scaling Params | Target Tokens | Auto Batch | |-------|---------------|---------------|------------| | d=8 | 42M | 0.44B | 2^18 = 262K | | d=10-16 | 70M-235M | 0.7B-2.5B | 2^19 = 524K | | d=18-26 | 324M-918M | 3.4B-9.6B | 2^20 = 1.05M | | d=32-50 | 1.7B-6.2B | 17.6B-65.6B | 2^21 = 2.1M | In particular, this matches empirical observations that d26 prefers ~2^20 while d12 prefers ~2^19. ### Code Cleanup Also refactored model initialization to use `build_model_meta(depth)` helper and `dataclasses.asdict()` for cleaner config handling. ### Useful references - [Bergsma et al., Power Laws for Batch Size, Model Size, and Training Horizon](https://arxiv.org/abs/2505.13738) - [McCandlish et al., An Empirical Model of Large-Batch Training](https://arxiv.org/abs/1812.06162) - [Brown et al., Language Models are Few-Shot Learners](https://arxiv.org/abs/2005.14165) - [Merrill et al., The Batch Size–Critical Batch Size Myth](https://arxiv.org/abs/2505.23971) ### One more thing (batch size ramp) Tried batch size ramping. The simplest implementation I could think of "tricks" the existing training loop by slicing each micro-batch into smaller pieces and calling optimizer.step() more frequently early in training (1/8 → 1/4 → 1/2 → full batch over the first x% of training, with sqrt LR scaling). Also required a torch.compile warmup phase to pre-compile all slice sizes and avoid recompilation spikes during training. While the idea is sound and small gains were observed, they weren't sufficient to justify the code complexity introduced (conditional slicing logic, warmup with state save/restore, etc.). Not merged for now. --- ## 2026-02-05: SwiGLU Activation (Negative Result) Replaced ReLU² MLP activation with SwiGLU (inspired by [twitter](https://x.com/_xjdr/status/2019141521690567058)). SwiGLU uses three projections instead of two, so to match parameters and FLOPs we scale hidden_dim from 4× to 8/3×: ```python # Old ReLU²: 2 matrices, 4x expansion # params: 2 × n × 4n = 8n² # flops: 2 × 2n × 4n = 16n² per token self.c_fc = Linear(n_embd, 4 * n_embd) self.c_proj = Linear(4 * n_embd, n_embd) x = c_proj(relu(c_fc(x)).square()) # New SwiGLU: 3 matrices, 8/3x expansion # params: 2 × n × (8n/3) + (8n/3) × n = 8n² ✓ matches # flops: 3 × 2n × (8n/3) = 16n² per token ✓ matches hidden_dim = (8 * n_embd) // 3 self.w1 = Linear(n_embd, hidden_dim) # gate self.w2 = Linear(n_embd, hidden_dim) # up self.w3 = Linear(hidden_dim, n_embd) # down x = w3(silu(w1(x)) * w2(x)) ``` Tested at both d12 and d24 (GPT-2 scale). Worse on all measures — step efficiency, wall clock time, and FLOPs. ReLU² remains superior for nanochat. **Not adopted.** --- ## 2026-02-03: Flip Muon MLP LR Multiplier (PR #492) Tested flipping the shape-based LR heuristic in Muon from boosting tall matrices (input projections like `c_fc`) to boosting wide matrices (output projections like `c_proj`). The original code applies `max(1, rows/cols)^0.5`, giving ~2x LR to `c_fc`. The flipped version gives ~2x LR to `c_proj` instead, which aligns with classical fan-in/fan-out scaling conventions. This was proposed in [PR #492](https://github.com/karpathy/nanochat/pull/492) and showed improvements in modded-nanogpt. **Result:** Quick d12 experiment: slightly worse **Not adopted.** --- ## 2026-02-03: Skip AdamW Every Other Step Inspired by modded-nanogpt, tried stepping AdamW only on odd iterations while Muon steps every iteration. The idea is that small AdamW params (embeddings, scalars, gates) don't need updates as frequently as the large weight matrices, and skipping saves both compute and communication. Added `skip_adamw` parameter to `MuonAdamW.step()` and `DistMuonAdamW.step()` plus a matching `zero_grad(skip_adamw=...)` to let AdamW gradients accumulate over 2 steps. Used `lr *= 2**-0.5` (sqrt scaling) to compensate for the 2x effective batch size on AdamW params. **Result:** for nanochat d12, we see ~2% faster tok/s, but each step is slightly worse in loss. On net, when plotting against wall clock time, it's slightly worse. **Not adopted.** --- ## 2026-02-02: FP8 Training with torchao Integrated FP8 training using `torchao.float8` to accelerate Linear layer matmuls on H100 GPUs. ### Background FP8 (8-bit floating point) uses H100's FP8 tensor cores for ~2x theoretical matmul throughput. The tradeoff is quantization overhead: computing scales and casting tensors to/from FP8. Still, as an example torchtitan (Meta's distributed training framework) reports 25-28% speedups with FP8 for some of their experiments. **Previous attempt (Jan 2026):** FP8 on just `lm_head` following modded-nanogpt with custom ops → 1% speedup, +2GB memory. Failed due to fragile torch.compile interaction. But this experiment was also done on ~d12 scale back then instead of the bigger model that gets GPT-2 capability of approx d24. **This attempt:** Use torchao's `convert_to_float8_training()` on ALL Linear layers, increase model size to d24. The core snippet is: ```python from torchao.float8 import Float8LinearConfig, convert_to_float8_training config = Float8LinearConfig.from_recipe_name("tensorwise") convert_to_float8_training(model, config=config) ``` But in practice it's more involved (see base_train.py). ### Results **Microbenchmark (d26 MLP, 65536x1664 @ 1664x6656):** | Method | Forward | Fwd+Bwd | Speedup | |--------|---------|---------|---------| | BF16 + compile | 2.00ms | 4.79ms | 1.00x | | FP8 rowwise + compile | 1.84ms | 4.55ms | 1.08x | | FP8 tensorwise + compile | 1.45ms | 4.06ms | **1.38x** | | FP8 rowwise (no compile) | 2.89ms | 21.86ms | 0.23x ❌ | torch.compile is MANDATORY. Without it, FP8 is 4x slower due to unfused scaling ops. **Full training (d26):** | Config | tok/sec | vs baseline | |--------|---------|-------------| | BF16 baseline | 630K | 1.00x | | FP8 rowwise | 564K | 0.90x ❌ | | FP8 tensorwise | 740K | **1.17x** ✓ | Memory usage also decreases quite a bit, by ~9GB (activations stored as FP8 instead of BF16). Seeing 17% speedup is encouraging but we're still not done yet because each step is now in lower precision and less powerful individually, so to make up for the precision drop we have to train longer. Empirically, running some sweeps overnight on d24 scale, I saw that the actual speedup (when you match performance) is closer to 5%. It's possible that our LLMs at ~d24 scale are still too small to confidently enjoy the speedups that come from fp8 for bigger models. ### Key Learnings For nanochat at approximate scale of interest (~GPT-2 capability, ~d24): 1. **Tensorwise >> Rowwise** - Rowwise computes per-row scales, overhead exceeds benefit. Tensorwise uses one scale per tensor. 2. **Filter small layers** - Layers with dims not divisible by 16 must be skipped (FP8 hardware requirement) 3. **Larger models benefit more** - d12 was still slower with FP8; d26+ shows gains. Therefore, in some depths there is a benefit to fp8 and in some there isn't. Keeping it configurable for now, passed in via kwargs and default off. 4. **The effective, capability-matched speedup is lower still** - because each step is of slightly lower precision/quality. ### Integration Added `--fp8` flag to `base_train.py`, default recipe is "tensorwise", example of turning on: ```bash torchrun --nproc_per_node=8 -m scripts.base_train --depth=24 --fp8 ``` Uses tensorwise by default. Requires `torchao==0.15.0` (compatible with torch 2.9.1), which was added to dependencies. **TLDR**: turning on fp8 for GPT-2 capability nanochat model gives approx +5% capability-matched speedup. --- ## 2026-01-29: Hyperball/MuonH Experiments (Negative Result) Explored Hyperball optimization from [this post](https://psychedelic-sunstone-851.notion.site/Fantastic-Pretraining-Optimizers-and-Where-to-Find-Them-2-1-Hyperball-Optimization-2e924306e6f280e7a5ffee00eb40a0dd) (saved to `knowledge/muonh.md`). Constrains weights to sphere of radius R (initial norm): `W_{t+1} = R · Normalize(W_t - η·R · Normalize(u_t))`. Had to change a number of details in a branch, e.g. not use zero init for our projections (or the initial norm would be zero), keep track of the initial norm, adjust Muon -> MuonH for the update. Experiments on d12: | Experiment | Result | |------------|--------| | MuonH for matrix params | Worse than baseline | | MuonH + LR sweep (2.5e-3 to 1e-2) | Still worse | | Added learnable RMSNorm scales (paper says γ preserves expressivity) | Still worse | | Various RMSNorm init tweaks, e.g. 0 at init to residual | Still worse | | AdamH for lm_head (paper recommends this) | Broken - loss plateaus (see below) | | AdamH + learnable output scales | Still worse | Could not outperform the baseline implementation. The article doesn't go into too much detail on how AdamH is applied to `lm_head` exactly. The classifier layer has to be able to increase in magnitude to make more confident predictions over time. Tried a sensible version with added 0-D learnable scalar, and also with RMSNorms with per-channel learnable scalars both pre and post resnet blocks. **Result:** This was not an out-of-the-box win for nanochat even with a mild attempt over a few hours at a bit of tuning and debugging. The idea itself is intuitively appealing. Might come back around later to try harder later. --- ## 2026-01-28: Reverted Bigram Hash Embeddings Removed bigram embeddings (engram-lite) from the codebase. At larger scale (d25), the improvement was tiny and disappeared entirely when measured by wall clock time. It also bloated the VRAM used. The extra parameters and complexity aren't justified. --- ## 2026-01-27: Bigram Hash Embeddings (Engram-lite) Explored N-gram memory modules inspired by the [DeepSeek Engram paper](https://arxiv.org/abs/2601.07372) and [modded-nanogpt PR #201](https://github.com/KellerJordan/modded-nanogpt/pull/201). ### Background The Engram paper introduces "conditional memory" as a complement to MoE - using O(1) hash lookups to retrieve static N-gram patterns instead of reconstructing them through computation. Key insight: transformers waste early layers "simulating retrieval through computation" for patterns like named entities and formulaic phrases that could be simple table lookups. ### What We Tried **1. Full Engram module with context-aware gating (paper design)** ```python # Hash bigrams to retrieve embeddings, then gate with hidden state e = embed(hash(prev_token, curr_token)) q = RMSNorm(h) # hidden state as query k = RMSNorm(W_k @ e) # projected embedding as key v = W_v @ e α = sigmoid(q · k / √d) # scalar gate per position output = α * v ``` - Injected after block 1 (paper found early injection optimal) - Slight improvement, but quite a bit of complexity added. **2. Early-layer only injection** - Only inject bigram signal in first 4 layers (where paper claims static pattern offloading helps most) - **Result:** Actually hurt performance. The model seems to need uniform injection across all layers. **3. Trigrams** - Extended to hash both 2-grams and 3-grams, concatenating embeddings - **Result:** No improvement over bigrams alone. Dilutes capacity from more frequent 2-gram patterns. **4. Bigram-only with x0-style injection (modded-nanogpt engram-lite approach)** - Simple hash: `(36313 * curr) XOR (27191 * prev) mod table_size` - Zero-init embedding table, learned per-layer lambdas - Add to residual at every layer: `x = resid_λ[i]*x + x0_λ[i]*x0 + bigram_λ[i]*x0_bigram` - **Result:** This simple approach works and provides a consistent improvement. TLDR The winning approach follows modded-nanogpt's "engram-lite", simply adding the following module and feeding its output into the residual branch (gated by a per-layer learnable \lambda) before every single block: ```python class BigramEmbed(nn.Module): def __init__(self, vocab_size, embed_dim, table_multiplier=5): self.embed = nn.Embedding(vocab_size * table_multiplier, embed_dim) def forward(self, idx): h = (36313 * idx[:, 1:]) ^ (27191 * idx[:, :-1]) % (table_size - 1) return self.embed(h) ``` As for optimal hyperparameters: - **Table size:** `vocab_size * 5` (~164K entries for 32K vocab). Swept a number of settings and 5 was optimal. - **Injection:** Every layer via learned `bigram_lambdas` (init 0.1 was better than 0.0). - **Normalization:** Also tried adding a `norm()` to the embeddings (mirroring the token embeddings), this was slightly worse. - **Init:** Zero-init embedding, so starts as identity (tried small noisy init, it's worse) - **Optimizer:** AdamW with same LR as token embeddings ### Key Learnings 1. **Gating didn't help at our scale.** The paper's context-aware gating mechanism (sigmoid dot-product gate) added parameters and complexity without improvement. modded-nanogpt found the same: "simple direct addition to the residual stream outperformed by a decent margin." 2. **Uniform injection beats early-only.** Despite the paper's finding that early layers benefit most, restricting injection to early layers hurt. The x0-style "add everywhere with learned lambda" pattern works better for our architecture/scale. 3. **Bigrams are sufficient.** Trigrams didn't help - the extra context doesn't pay for the diluted capacity. 4. **Scale matters.** The Engram paper's results are at 27B params with MoE. At our ~100M-1B scale, the simpler approach wins. The elaborate gating mechanism may become useful at larger scales where collision handling matters more. ### Parameters Added For d12 model with `table_multiplier=5`: - Bigram embedding: 32768 × 5 × 768 = ~126M params - Per-layer lambdas: 12 scalars (negligible) If you're keeping track, we now have *a lot* of parameters, a significant amount of them in embeddings (token embeddings, bigram embeddings, value embeddings). For example, for a d12 we now have: ``` Parameter counts: wte : 25,165,824 bigram_embed : 125,829,120 value_embeds : 150,994,944 lm_head : 25,165,824 transformer_matrices : 84,935,808 scalars : 36 total : 412,091,556 ``` In other words, only about a quarter of parameters are now weight projections and the vast majority is embedding tables. Still, on all axes (steps, wall clock time, flops), this somewhat parameter-bloated architecture beats the baseline and will now become the default. After adding the engram-lite, I re-ran the scaling laws to determine the new optimal tokens:params ratio. I swept FLOPs in the range 1e18..1e19, exponentially strided in 4 settings (1e18, 2e18, 5e18, 1e19). I looked at a number of ways of determining the effective parameter count for the purposes of the scaling laws. The results looked like this: ``` Kaplan-style (all projections including lm_head and no embeddings) Optimal configurations (from quadratic fits): FLOPs Eff Params Tokens Ratio Val BPB ----------------------------------------------------------------- 1e+18 110,678,115 1,241,505,403 11.2 0.8972 2e+18 167,797,457 1,785,336,422 10.7 0.8616 5e+18 250,650,865 2,642,234,152 10.8 0.8293 1e+19 381,758,347 3,806,871,243 10.3 0.7999 N \propto C^0.54, D \propto C^0.49 Chinchilla-style (all parameters, period.) Optimal configurations (from quadratic fits): FLOPs Eff Params Tokens Ratio Val BPB ----------------------------------------------------------------- 1e+18 416,320,605 1,232,157,011 3.0 0.8974 2e+18 560,239,841 1,763,669,281 3.2 0.8616 5e+18 741,495,903 2,629,909,368 3.6 0.8291 1e+19 988,644,331 3,884,841,895 4.0 0.7999 N \propto C^0.37, D \propto C^0.50 Transformer-only-style (only the projections inside the transformer) Optimal configurations (from quadratic fits): FLOPs Eff Params Tokens Ratio Val BPB ----------------------------------------------------------------- 1e+18 80,259,665 1,315,639,547 17.2 0.8966 2e+18 131,488,566 1,864,134,141 14.5 0.8622 5e+18 220,985,474 2,595,328,843 12.1 0.8302 1e+19 401,213,504 3,328,704,512 8.5 0.7994 N \propto C^0.70, D \propto C^0.41 ``` Clearly, the Kaplan-style ratios are most consistent and produce stable ~0.5 exponents for both params and tokens, meaning we can have a single fixed ratio of tokens:params for compute optimal models. This turns out to be about ~10.5, which now becomes the new default. --- ## 2026-01-19 to 2026-01-22: Optimizer Hyperparameter Sweep Ran ~320 experiments across 6 rounds, scaling from d12→d16→d20 to find optimal optimizer hyperparameters. Added granular per-component control to `setup_optimizers()` — separate LRs and betas for embedding, unembedding, value_embeds, resid_lambdas, x0_lambdas, and Muon matrix params. ### What We Swept - Learning rates for all 6 parameter groups - Beta1/beta2 for all 5 AdamW groups - Muon momentum (start/end), weight decay - Hundreds of combinations (2-way, 3-way, 4-way, etc.) ### The Journey **At d12**, found two independent improvement routes: - **Route A:** emb_lr↑ (0.3→0.4), weight_decay↑ (0.1→0.15), matrix_lr↑ (0.02→0.025) - **Route B:** x0_lr↓ (0.5→0.2), x0_beta1↑ (0.8→0.9+) Both gave ~0.002 improvement, but combining them caused conflicts. Fine-tuning found wd=0.13, matrix_lr=0.027, emb_lr=0.38 helped slightly. Best d12 config: Route A + x0_beta1=0.95. **At d16**, Route B became competitive with Route A. The routes still conflicted when combined. **At d20** (target scale), everything changed: - Fine-tuned values from d12 **actively hurt** performance - Routes no longer conflicted - Just `x0_beta1=0.96` alone captured nearly all the gains ### Final x0_beta1 Sweep at d20 | x0_beta1 | val/bpb | Δ vs baseline | |----------|---------|---------------| | **0.96** | **0.7971** | **-0.0007** | | 0.94 | 0.7972 | -0.0006 | | 0.90 | 0.7972 | -0.0006 | | 0.97 | 0.7977 | -0.0001 | | 0.98 | 0.8011 | +0.0033 💀 | Flat plateau from 0.90-0.96, then sharp cliff at 0.97+. ### Key Learnings 1. **Hyperparameters are scale-dependent.** What works at d12 doesn't transfer to d20. The elaborate fine-tuning that won at d12 actively hurts at d20. 2. **Improvement magnitude shrinks with scale.** ~0.002 at d12 → ~0.0007 at d20. The baseline is already better-tuned for larger models. 3. **Sharp cliffs exist.** x0_beta1=0.98 is catastrophic while 0.96 is optimal. 4. **Don't over-tune on small proxies.** Validate at target scale before shipping. ### Final Recommendation For production d20 runs, add one flag: ``` --x0-lambdas-beta1=0.96 ``` Skip everything else discovered at smaller scales. --- ## 2026-01-18: More various experiments - Tried Muon custom kernels for XXT and all the others. The improvement was there for targeted tests (~20%) but washed out completely to noise in an actual training run, especially because the Muon compute is split across all the workers. Abandoned due to complexity bloat. - Fuse Q,K,V,O nn.Linear layers into a single QKVO Linear layer. ~Zero impact - Tried the `sa_lambdas` that gate QKV and O. Slightly confused because of the use of rmsnorm, which erases the effect of any scalar multiplier. Helped a tiny bit (~1e-4 of loss), abandoned to control complexity. --- ## 2026-01-17: Various experiments Modded-nanogpt uses [Value Embeddings](https://arxiv.org/abs/2410.17897) (VEs) in a funny U-shaped structure, 3 of them in total and with gates. I tried a large number of tweaks on this today: - VEs at every layer, at alternating layers, U shaped, front and back. Alternating layers worked best, i.e. we end up with *a lot* more VEs than modded-nanogpt, at every other layer. It works better. - Many parameters sharing ideas to reduce new parameter count, nothing here worked. All failed. - Many ideas to reduce parameter count, the LLM hates all of them: low rank decompositions, projections. All failed. - Gated yes or no and how much. Gate helps. Long story short is that the models *love* Value Embeddings. It is a way to add a huge amount of capacity (parameters) to the model at almost zero cost of FLOPs, because these embeddings are simply added to the Values tensor. Any attempt to reduce the capacity of value embeddings (param sharing, low rank, projections) fail. The model wants many of them, and with all the capacity, and doing so wins across all x axes of steps, flops and wall clock. I re-ran the scaling laws and, because the models are now very parameter bloated, the optimal ratio has halved from 8 to 4! Way down lower than Chinchilla's 20 at this point. Other experiments, looking at val/bpb as a function of all of steps, flops and wall clock time: - Aspect ratio of 128 is worse than 64, I tried a sweep fixing FLOPs == 1e18 and 64 outperforms. The LLM prefers to be slightly thinner and longer. - Head dim definitely prefers to be 128 instead of 64, i.e. fewer bigger heads - Bunch of other random stuff like that. Keeping all of this work on a private branch for now but hope to push shortly. --- ## 2026-01-17: Modded-nanogpt Ideas Sweep (Continued) Continued testing ideas from modded-nanogpt. | Idea | Result | Notes | |------|--------|-------| | Attention gates | No improvement | Per-head learnable gates on attention output. +1GB memory, decreased efficiency. | | Batch size schedule | Abandoned | 8→16→24 with LR scaling. Made training script too bloated/complex, not worth cognitive overhead. | | Value embeddings | Helps a lot | Experiments still ongoing, more on this later. | --- ## 2026-01-16: Flash Attention 3 Fallback to SDPA Added automatic fallback from Flash Attention 3 to PyTorch's `scaled_dot_product_attention` (SDPA) for users without Hopper GPUs. This enables nanochat to run on older CUDA GPUs, CPU, and MPS (Apple Silicon). ### Implementation Created `nanochat/flash_attention.py` - a unified interface that: - Detects FA3 availability at import time (requires sm90+ / Hopper) - Exports a `flash_attn` object matching FA3's API exactly (`flash_attn.flash_attn_func`, `flash_attn.flash_attn_with_kvcache`) - Automatically routes to FA3 or SDPA based on hardware - Handles tensor layout differences: FA3 uses (B, T, H, D), SDPA uses (B, H, T, D) - Implements sliding window attention via explicit masks for SDPA - Manages KV cache manually for SDPA (FA3 does it in-place) ### Changes to Existing Files Changes to existing code were intentionally kept extremely minimal. **gpt.py**: Only the import line changed and a comment **engine.py**: Zero changes needed **base_train.py**: Added status print and warnings: - Prints whether FA3 or SDPA fallback is being used - Warns about efficiency loss without FA3 - Warns about sliding window support if `--window-pattern` is not "L" ### Testing Tests are split into two classes due to dtype/device constraints: 1. **TestFA3VsSDPA**: Comparison tests requiring Hopper GPU + bfloat16. Run both implementations on identical inputs and verify outputs match (max diff typically 0, at most ~0.004 for sliding window). 2. **TestSDPAOnly**: SDPA-only tests that run on any device with appropriate dtype. Verify forward pass, backward pass, and KV cache work correctly. Added `_override_impl` mechanism for testing - can force 'fa3' or 'sdpa' to directly compare implementations. ### Notes - SDPA fallback is significantly slower than FA3 especially in that it lacks the sliding window attention support - Recommend `--window-pattern L` (full context) when using SDPA fallback --- ## 2026-01-16: Modded-nanogpt Ideas Sweep (Mostly Negative) Tested several architectural ideas from modded-nanogpt to see if they transfer to nanochat. All of these did not help: | Idea | Result | Notes | |------|--------|-------| | Half-truncated RoPE | No improvement | Only first half of head dims get RoPE (base 1024, linspace). Second half "stationary". | | Asymmetric softcap | Slightly worse | `23 * sigmoid((x+5)/7.5)` vs our symmetric `15 * tanh(x/15)`. May only help with FP8. | | Smear gate | Negligible | Blend each token with predecessor via learned gate. Tiny improvement not worth n_embd² params. | | Backout | No improvement | Save activations at ~60% through network, subtract scaled version at end. | | Skip connection | Slightly worse | Save at layer ~25%, add at layer ~50%. Also +2GB memory from storing activations. | Value Embeddings do show promise. I need a more elaborate exploration of a few related ideas, which I leave for tomorrow. --- ## 2026-01-15: Olmo pretraining mix (Negative result) I attempted to train on the Olmo 3 pretraining dataset [allenai/dolma3_mix-6T](https://huggingface.co/datasets/allenai/dolma3_mix-6T) instead of FineWeb-edu. I ran into a number of [errors and issues](https://huggingface.co/datasets/allenai/dolma3_mix-6T/discussions/2) trying to both download and process the dataset and then noticed some quality issues (e.g. some documents seem to be extremely short, like "5".). I managed to work around these with some sensible hacks (e.g. reject documents less than 100 characters in length) and tried to process the dataset exactly as FineWeb, re-trained the tokenizer and trained a d16 model. The CORE score decreased from 15.5 to 13.8, i.e. the result is quite a bit worse. I am still looking to try the [DCLM dataset](https://arxiv.org/abs/2406.11794), which according to the paper should be better that FineWeb-edu. I do have some concerns that the same group both prepared the DCLM dataset *and* introduced the CORE score so I'm a bit hesitant in case there was some overfitting to CORE score adjacent data distribution. Classifying as negative result and reverting back to FineWeb-edu for now. --- ## 2026-01-13: Varlen Attention (Negative Result) Attempted to prevent attention from "leaking" across document boundaries using Flash Attention's `flash_attn_varlen_func`, similar to modded-nanogpt's approach. ### Background With the BOS-aligned dataloader, multiple documents are packed into each row. Standard attention allows tokens to attend across document boundaries within a row. The hypothesis was that preventing this "leakage" via varlen attention might improve training. ### Approach: Compute cu_seqlens from inputs - Find BOS positions: `(inputs.view(-1) == bos_token_id).nonzero()` - Gotcha 1: Variable-length `cu_seqlens` caused torch.compile recompilation (25s/iter!) - fixed by padding to fixed size - Gotcha 2: `nonzero()` inside compiled model hit recompile limit - fixed by moving computation outside compiled region ### Final Results (d16) | Metric | Baseline | Varlen | |--------|----------|--------| | val_bpb | 0.85427 | 0.85407 | | MFU | ~same | ~same | | tok/sec | ~same | ~same | Essentially identical. The 0.0002 bpb improvement is almost noise. ### Conclusion Not worth the code complexity. The "leakage" across document boundaries within a row is not harmful - the model handles it fine. The BOS-aligned dataloader already provides the key benefit (every row starts with proper context). Not merging to master. --- ## 2026-01-13: BOS-Aligned Dataloader with Bin Packing Redesigned the pretraining and midtraining dataloader to ensure every sequence starts with a BOS token, and explored bin-packing algorithms to minimize wasted tokens. ### Problem Statement The original dataloader streams tokens into a flat buffer and reshapes into batches. This means some rows start mid-document (no BOS), which could confuse the model during training. We want every row to start with BOS and contain well-formed documents. ### Approach 1: Greedy-Crop BOS (Simple) Each row is built independently: - Start with a document (which has BOS prepended) - Pack more documents until row is full - If a document doesn't fit, **crop it** to fill remaining space (discard the rest) - 100% utilization (no padding), but wastes cropped tokens ### Waste Analysis Measured token waste empirically on real data (T=2048): - **39.4% of tokens are cropped** (discarded when docs don't fit) - **22.9% is the theoretical minimum** (tokens in docs longer than T+1 that can never fit) - The extra ~16.5% comes from "unlucky" cropping when a long doc starts near the end of a row ### Bin Packing Algorithms Explored | Algorithm | Util% | Crop% | Pad% | Notes | |-----------|-------|-------|------|-------| | Greedy-Crop (baseline) | 100% | 39.4% | 0% | Simple, no wasted compute | | Greedy-Pad | 78% | 23.0% | 22% | Pads instead of crops - wastes compute | | First-Fit Decreasing (FFD) | 99.7% | 23.0% | 0.3% | Near-optimal packing, minimal padding | | **BestFit-Crop** | 100% | 34.6% | 0% | Smart cropping, no padding | ### BestFit-Crop Algorithm A middle ground that maintains 100% utilization while reducing cropping: 1. Buffer N documents 2. For each row, greedily pick the **largest doc that fits entirely** 3. Repeat until nothing fits 4. When nothing fits, crop a doc to fill remaining space exactly This avoids "unlucky" crops by searching the buffer for better-fitting documents. **Results (T=2048):** - Crop waste reduced from 39.4% → 34.6% (~12% relative improvement) - Still achieves 100% utilization (no padding, every token trains) - Slightly more rows than baseline (uses more documents per batch) ### Decision: Keep Two Implementations 1. Keep the original implementation which is very simple, efficient and has 100% token utilization in the batch (no padding with ignore tokens), but creates slightly more confusing token streams for the LLM because documents during training can start abruptly from the middle with no context. Note that this never happens at test time, where BOS is always present. 2. **`_bos_bestfit` (BestFit-Crop, new default)**: Slightly more complex but still keeps 100% token utilization in the batch (no padding), but at the cost of discarding documents when they don't fit. In practice, about 34% of tokens are discarded with this approach. This is ok because for most models we care about we have plenty of data without having to go to multiple epochs. One more subtle effect is that it does skew the data distribution a tiny bit because, reliably and necessarily, tokens at the tails of long documents will be discarded. However, this doesn't seem to impact actual downstream performance. ### Midtraining The midtraining dataloader was also updated. Because conversations are on average a lot shorter than pretraining documents, only about 3.3% of tokens get cropped. ### NOTE: loss scale Do note that switching to the BOS dataloader changes the validation loss and makes all previous experiments not comparable in absolute value of the loss, because we have a lot fewer "confusing" tokens in the train/val batches. All tokens can look back and find the BOS token and have the full context of that document to make predictions. Therefore, the loss appears lower but this is "fake" to some extent, and the expectation is that the vast majority of relative comparisons done so far would agree with those before and after this change. --- ## 2026-01-13: Number Token Split Pattern Validated the `\p{N}{1,2}` pattern in `SPLIT_PATTERN` (tokenizer.py line 30), which I only guessed earlier and had a TODO for to validate. GPT-4 uses `\p{N}{1,3}` to group number sequences of up to 3 digits into tokens, but we suspected smaller vocab sizes benefit from grouping fewer digits per token. **Results (d12, vocab=32K):** | Pattern | val_bpb | |---------|---------| | `\p{N}{1,1}` | 0.969 | | `\p{N}{1,2}` | **0.965** | | `\p{N}{1,3}` | 0.972 | **Conclusion:** `{1,2}` is optimal for vocab size 32K. Grouping 3 digits wastes tokens on rare 3-digit combinations; grouping 1 digit is too fine-grained and bloats token sequences. Keeping `{1,2}` as default. --- ## 2026-01-13: FP8 Training for lm_head Attempted to use FP8 (8-bit floating point) for the lm_head layer to speed up the large vocab projection matmul. H100 GPUs have FP8 tensor cores that can theoretically provide ~2x speedup over BF16. ### Implementation Approaches Tried **1. Dynamic Scaling (failed)** - Compute `x.abs().max()` and `w.abs().max()` each forward to determine scales - Problem: `.item()` calls cause graph breaks with torch.compile - Tried `@torch._dynamo.allow_in_graph` pattern (like torchao.float8) - worked but no speedup - Tried `torch.library.custom_op` with float scales - caused NaN gradients after first optimizer step - Root cause: interaction between custom ops, dynamic scale computation, and torch.compile is fragile **2. Static Scaling (partial success)** - Pre-set scales at init time like modded-nanogpt: `x_scale=10/448, w_scale=0.1/448` - `grad_scale` computed dynamically from batch size (safe since it's just `1/(B*T)/57344` due to the gradient expression of cross entropy). modded-nanogpt has a bug here probably because they set `grad_scale = 0.75/448`, but grads are in E5M2 so this should probably be `1/57344`, 1 being the amax of any individual element of cross entropy loss, and no normalization by B,T because they use sum reduction not mean reduction. - Uses `torch.library.custom_op` with `@torch.compile` on inner kernels - This works correctly - no NaNs, proper gradients ### Results (d12) | Metric | BF16 Baseline | FP8 lm_head | |--------|---------------|-------------| | GPU Memory | 34 GB | 36 GB | | tok/sec | baseline | ~1% faster | ### The Memory Mystery FP8 *should* save memory since we store `x_f8` (1 byte) instead of `x` (2 bytes) for backward. But we see 2GB *increase*. Suspected causes: - `torch.compile` on inner kernels creating extra buffers/specializations - `torch._scaled_mm` internal workspace allocations - Custom op registration machinery overhead Tried saving original weight `w` (just a reference to parameter) instead of `w_f8` in backward, then re-quantizing on the spot during backward - didn't help. Still saw bump. ### Microbenchmark vs Reality Raw microbenchmark showed promise: - BF16 matmul: 16.95 ms - FP8 matmul (static scales): 10.31 ms (1.64x faster) - FP8 with dynamic scaling: 12.25 ms (1.38x faster) But in full training, the ~1% tok/sec improvement doesn't justify the 2GB memory increase and the added code complexity and the need to tune scale factors for both x and w. ### Code Artifacts See the branch `fp8_attempt_fail` for: - `nanochat/fp8_static.py` - Static scaling implementation (working) - `nanochat/fp8_dynamic.py` - Dynamic scaling implementation (torchao-style, working but slow) - `gpt.py` imports `fp8_static.LinearFP8` and simply swaps it for `lm_head` in `gpt.py`. ### Open Questions - Why does the custom op approach use more memory than vanilla BF16? - Why is the bump in tok_per_sec so low? We should see ~1.6X speedup in both the forward pass and also (twice) in backward pass for the gradients. Granted, Amdahl's law is part of the solution because our vocab_size is only 32K so the final layer isn't a huge part of the profile but the expected speedup is still not fully realized. **Conclusion:** Negative result for now. The implementation works correctly but provides marginal speedup with *increased* memory usage. I'm not understanding the torch.compile interaction here. The complexity of FP8 custom ops isn't justified for lm_head alone. TODO to study in more detail the way this is implemented in other libraries, e.g. torchao. --- ## 2026-01-12: Multi-Token Prediction (MTP) Ported multi-token prediction from modded-nanogpt. Instead of predicting just the next token, predict the next n tokens at each position with weighted loss. ### Implementation - Instead of calling the loss `n_predict` times, uses a fancy batched computation using `unfold` + `gather` + cross-entropy decomposition (`CE = logsumexp - logits[target]`) - Schedule anneals from 3-token to 1-token prediction: - 0-33%: `[1.0, 0.5, 0.25→0]` (3rd token fades) - 33-67%: `[1.0, 0.5→0]` (2nd token fades) - 67-100%: `[1.0]` (standard next-token) - Weights normalized to sum to 1 ### Results (d12) | Metric | Baseline | MTP | |--------|----------|-----| | GPU Memory | 34 GB | 47 GB | | MFU | 41% | 40% | | val/bpb (per step) | baseline | same/slightly worse | | val/bpb (wall clock) | baseline | noticeably worse | **Conclusion:** Negative result for nanochat. The extra memory and compute overhead from predicting multiple tokens doesn't pay off, in fact the results get worse. The auxiliary loss signal may help in other settings (larger models, different architectures?), but for our setup it's pure overhead at the moment. --- ## 2026-01-11: Sliding Window Attention Added configurable sliding window attention, inspired by GPT-3's alternating short/long pattern. **Pattern string configuration:** - New `--window_pattern` CLI arg and `GPTConfig.window_pattern` field - Pattern is tiled across layers (e.g., `SSSL` for 20 layers → `SSSLSSSLSSSLSSSLSSSL`) - Final layer always forced to L (full context) regardless of pattern - Short window = `sequence_len // 2` - Long window = `sequence_len` (full context) - All previous models so far have been simply `L` and checkpoint loading is modified accordingly to fill in this param for old models, see `_patch_missing_config_keys` Quick experiments showed `SSSL` (every 4th layer is long) works well - provides a good balance between compute savings and model quality. This is now the default. --- ## 2026-01-11: Flash Attention 3 Integration Replaced PyTorch's `scaled_dot_product_attention` (FA2) with Flash Attention 3 for training and inference. ### Changes Made **1. FA3 via `kernels` package** - Official FA3 is "beta" and requires building from source (painful) - Using `kernels` package from HuggingFace Hub: `get_kernel('varunneal/flash-attention-3')` - Loads pre-built wheels, works out of the box on H100 **2. Simplified attention code** - FA3 uses `(B, T, H, D)` layout matching our projection output directly - no transpose needed - Training: `flash_attn.flash_attn_func(q, k, v, causal=True)` - Inference: `flash_attn.flash_attn_with_kvcache()` handles all cache cases in one call - Removed 3 separate FA2 code paths (training, single-token, chunk inference) - GQA handled automatically when n_kv_heads < n_heads **3. Rewrote KVCache for FA3** - Old format: `(num_layers, 2, B, H, T, D)` combined tensor - New format: separate `k_cache` and `v_cache` of shape `(num_layers, B, T, H, D)` - FA3 updates cache in-place during `flash_attn_with_kvcache` - Position tracked via `cache_seqlens` tensor (int32, per batch element) - Simpler API: `get_layer_cache()`, `advance()`, `reset()`, `prefill()` ### Results - **~9% improvement in tok/sec** during training out of the box - Benchmarks showed FA3 is 2x faster than FA2 at realistic training sizes (batch=32, seq=2048) - FA3 supports sliding window via `window_size=(left, 0)`, which is huge and expected to give further improvements. This is ready to tune but keeping full context for now. --- ## 2026-01-11: Per-Layer Residual Scalars (x0 & resid lambdas) Cherry-picked an idea from modded-nanogpt around learnable per-layer residual connections. ### Changes Made **1. x0_lambdas (x0 residual connections)** - Save initial normalized embedding as `x0` after `norm(wte(idx))` - At each layer, blend x0 back in: `x = resid_lambdas[i] * x + x0_lambdas[i] * x0` - Zero-initialized, so disabled at start; model learns which layers benefit from the shortcut - Provides direct path from embedding to deep layers, helps preserve token information **2. resid_lambdas (residual stream scaling)** - Per-layer multiplicative scaling of the residual stream - Initialized to 1.0 (neutral, standard transformer behavior) - Allows model to learn to amplify/dampen residual at each layer **3. DistAdamW small parameter handling** - Added support for parameters with < 1024 elements (like the scalar lambdas) - Small params use `all_reduce` instead of `reduce_scatter`/`all_gather` - Fixes crash when param shape isn't divisible by world_size ### Key Finding: Different LR Sensitivity The two scalar types need very different learning rates: - **x0_lambdas (additive)**: Can use normal LR (~0.5). Adding a fraction of x0 is forgiving. - **resid_lambdas (multiplicative)**: Needs ~100x smaller LR (~0.005). Multiplying the residual compounds through layers. Implementation: `resid_params` gets `scalar_lr * 0.01`, `x0_params` gets full `scalar_lr`. ### Experiment Results Swept `--scalar_lr` (controlling x0_lambdas) at multiple depths: | Depth | Baseline (disabled) | Best scalar_lr | Best val_bpb | Δ bpb | |-------|---------------------|----------------|--------------|-------| | d8 | 1.0885 | 0.20 | 1.0782 | -0.0103 | | d12 | 0.9770 | 0.60 | 0.9693 | -0.0077 | | d16 | 0.9059 | 0.20 | 0.9002 | -0.0057 | | d20 | 0.8565 | 0.10 | 0.8526 | -0.0039 | **Observations:** - Consistent improvement across all model sizes - Optimal LR varies by depth; default of 0.5 is reasonable, but 0.6 is better for d12 - Adding resid_lambdas (with 0.01x LR) gives small additional improvement over x0 alone ### Meta Device Footgun Important lesson: `__init__` runs in meta device context, so any tensor values set there are fake. Must initialize actual values in `init_weights()`. Added docstring warning to `__init__`. ### Summary Added `--scalar_lr` (default 0.5) controlling learnable per-layer scalars. The formula `x = resid_lambdas[i] * x + x0_lambdas[i] * x0` gives the model control over residual scaling and direct shortcuts to the initial embedding. Solid improvement with essentially no compute overhead. --- ## 2026-01-10: Muon Optimizer Upgrades & Cautious Weight Decay Cherry-picked improvements from NorMuon (modded-nanogpt) into our simpler Muon implementation. Decided against using NorMuon directly due to hard-coded architecture assumptions (expects 32 params split 10 attn + 22 mlp), parameter labeling requirements, and complexity. ### Changes Made **1. Polar Express Orthogonalization** - Replaced Newton-Schulz iteration with "Polar Express Sign Method" from [arxiv.org/pdf/2505.16932](https://arxiv.org/pdf/2505.16932) - Uses 5 different coefficient tuples (one per iteration) instead of fixed coefficients - Both methods kept in code for easy comparison (`zeropower_via_polar_express` vs `zeropower_via_newtonschulz5`) - **Result:** No dramatic/noticeable difference in training, but keeping the new Polar Express as default. **2. NorMuon Variance Reduction** - Added per-neuron/column adaptive learning rate from NorMuon ([arxiv.org/pdf/2510.05491](https://arxiv.org/pdf/2510.05491)) - Maintains `second_momentum_buffer` with shape `[rows, 1]` or `[1, cols]` (whichever is smaller) - Normalizes updates based on running per-row/col variance estimate (beta2=0.95) - Memory overhead: ~1/max(rows, cols) per param, negligible - **Result:** Led to a very small improvement, kept and enabled by default. **3. Cautious Weight Decay** - Only decays weights where `update * weight >= 0` (same sign) from [arxiv.org/abs/2411.16085](https://arxiv.org/abs/2411.16085) - Standard WD always pulls toward zero; cautious WD skips decay when gradient is pushing weight away from zero - **Implementation note:** Had to inline the logic rather than use a separate `@torch.compile` function. Passing changing float values (like `weight_decay` during scheduling) as function arguments triggers recompilation. Reading from `group["weight_decay"]` inside the step avoids this. - **Result:** Solid improvements, especially the cautious version was better than standard wd. - Now defaults to ON for Muon via the `weight_decay` param. AdamW still has no weight decay and is hardcoded to 0 weight decay, might try to re-tune this later. **4. Weight decay schedule** - Added a linear schedule to weight decay that is default on from 1.0 to 0.0 (i.e. start with max weight decay in the beginning of training, then ramp to 0 by the end). Worked better than a static setting in experiments. (modded-nanogpt has the same schedule but it is implemented in a more confusing way by multiplying twice by the learning rate, which is already wired up to a decay schedule). ### Weight Decay Scaling Experiments Swept weight decay values at d8, d12, d16, d20 to find optimal values and scaling law. **Optimal Values Found:** | Depth | Width (channels) | Optimal WD | |-------|------------------|------------| | d8 | 512 | ~0.40 | | d12 | 768 | ~0.22 | | d16 | 1024 | ~0.10 | | d20 | 1280 | ~0.08 | **Scaling Law:** - Fit power law: `WD = k / channels^α` in log-log space - Found α ≈ 1.97 (approximately 2), meaning WD ∝ 1/width² **Practical Formula:** ``` WD_target = WD_reference × (d_reference / d_target)² ``` Example: If d12 optimal is 0.22, then d20 optimal ≈ 0.22 × (12/20)² ≈ 0.08 **Reference:** Moonlight paper uses fixed WD=0.1 for their 15B MoE model. Our experiments indicated a scaling law where the optimal WD changed with depth, so we go along with the empirical scaling law. ### Summary Muon was changed to use Polar Express, added NorMuon variance reduction, and cautious weight decay with schedule that ramps linearly to zero. All of these changes follow modded-nanogpt repo, but all of them were also validated piece by piece to yield improvements in nanochat with the exception of the Polar Express change which was in the noise. This is default on and configurable with `--weight_decay`, using simply 0.2 and ∝ 1/width² scaling. The kwarg `--weight_decay` is therefore changing as of this change. It used to configure AdamW via standard weight decay and now it becomes exclusively used in Muon (AdamW is hardcoded to 0.0), and it is scaled based on depth. --- ## 2026-01-08: exp_grad_clip - Gradient Clipping **Hypothesis:** Gradient clipping may be unnecessary overhead. Tested L2 norm clipping at various thresholds (0.25, 0.5, 1.0, 2.0) and elementwise clipping. **Results:** - No benefit at any scale tested (d12, d20) - All variants within noise (~0.9827 val_bpb) - Grad norm never exceeds 1.0 naturally, so clipping is always inactive - Clipping adds ~2% time overhead from the all-reduce **Bug Found:** Original implementation clipped local gradients before sync. Since this codebase doesn't use DDP (gradient sync is in the optimizers), each rank was clipping based on its own local norm. Fixed on the branch with proper distributed all-reduce. **Observation:** modded-nanogpt does not appear to clip either right now. **Summary:** Deleted all grad-clip code paths. The code naturally produces well-behaved gradients. This improves a bit of MFU because we don't have to calculate and sync grad norms. ================================================ FILE: dev/estimate_gpt3_core.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Estimating CORE Metric for GPT-3 Models\n", "\n", "**Authors**: Claude Code Opus 4.5, Andrej Karpathy\n", "\n", "**Date**: Jan 2026\n", "\n", "## Motivation\n", "\n", "The [CORE metric](https://arxiv.org/abs/2406.11794) (introduced in the DCLM paper) is a composite benchmark that evaluates pretrained language models across 22 diverse tasks spanning world knowledge, language understanding, commonsense reasoning, symbolic problem solving, and reading comprehension. It provides a single score that captures a model's general capabilities.\n", "\n", "We want to compare nanochat models against the GPT-3 model family from OpenAI's [\"Language Models are Few-Shot Learners\"](https://arxiv.org/abs/2005.14165) paper (2020). However, there's a problem: **GPT-3 models were never evaluated on CORE** (which didn't exist in 2020), and the models were never publicly released, so we can't evaluate them ourselves.\n", "\n", "## Our Approach\n", "\n", "We estimate CORE scores for GPT-3 by:\n", "\n", "1. **Identifying overlapping tasks** between the GPT-3 paper and CORE that were evaluated with similar methodology\n", "2. **Using GPT-2 as calibration data** — we have actual CORE scores for all 4 GPT-2 models, plus the GPT-3 paper reports results on GPT-2-equivalent tasks\n", "3. **Fitting a regression model** from the overlapping task scores to the full CORE score\n", "4. **Applying the model to GPT-3** using their reported task scores\n", "\n", "This notebook documents our methodology in detail for reproducibility." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "from pathlib import Path\n", "import pandas as pd\n", "\n", "# For nice table display\n", "pd.set_option('display.precision', 4)\n", "pd.set_option('display.max_columns', 20)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 1: Understanding CORE\n", "\n", "CORE consists of **22 tasks** evaluated in specific few-shot settings. The key innovation is **centering**: raw accuracies are adjusted to account for random guessing baselines.\n", "\n", "$$\\text{centered accuracy} = \\frac{\\text{accuracy} - \\text{baseline}}{1 - \\text{baseline}}$$\n", "\n", "The final CORE score is simply the **mean of all 22 centered accuracies**.\n", "\n", "### CORE Tasks\n", "\n", "| Category | Tasks |\n", "|----------|-------|\n", "| World Knowledge | Jeopardy, ARC Easy, ARC Challenge, BigBench QA Wikidata |\n", "| Language Understanding | HellaSwag (0-shot & 10-shot), LAMBADA, Winograd, Winogrande, BigBench Language ID |\n", "| Commonsense Reasoning | COPA, CommonsenseQA, PIQA, OpenBookQA |\n", "| Symbolic Problem Solving | BigBench Dyck, Operators, CS Algorithms, Repeat Copy Logic, AGI Eval LSAT-AR |\n", "| Reading Comprehension | SQuAD, CoQA, BoolQ |" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 2: Task Overlap Analysis\n", "\n", "We carefully compared the evaluation methodology between GPT-3 and CORE for each task. Key considerations:\n", "\n", "1. **Number of few-shot examples (K)**: GPT-3 often uses more examples than CORE\n", "2. **Task format**: Some tasks use different prompting strategies\n", "3. **Scoring method**: GPT-3 uses unconditional probability normalization for some tasks\n", "4. **Data split**: dev vs test set\n", "\n", "### Selection Criteria\n", "\n", "We applied a conservative filter: **both evaluations must use K=0 (zero-shot) or both must use K>0 (few-shot)**. We excluded tasks that mix zero-shot with few-shot, as this introduces systematic differences.\n", "\n", "### Tasks We Excluded\n", "\n", "| Task | GPT-3 K | CORE K | Reason for Exclusion |\n", "|------|---------|--------|----------------------|\n", "| Winograd | 7 | 0 | Mixing K>0 with K=0 |\n", "| Winogrande | 50 | 0 | Mixing K>0 with K=0 |\n", "| COPA | 32 | 0 | Mixing K>0 with K=0 |\n", "| OpenBookQA | 100 | 0 | Mixing K>0 with K=0, also uses unconditional normalization |\n", "| BoolQ | 32 | 10 | High sensitivity to K (17% gap between 0-shot and few-shot in GPT-3) |\n", "| CoQA | 5 | 0 | Different metric (F1 vs accuracy) |\n", "| LAMBADA few-shot | 15 | 0 | GPT-3 uses special fill-in-blank format |\n", "\n", "### Tasks Not in GPT-3 Paper\n", "\n", "These CORE tasks simply don't appear in GPT-3 (many didn't exist in 2020):\n", "- All 6 BigBench tasks (Dyck, Operators, CS Algorithms, Repeat Copy Logic, Language ID, QA Wikidata)\n", "- Jeopardy, CommonsenseQA, AGI Eval LSAT-AR\n", "- SQuAD v1 (GPT-3 uses v2)\n", "\n", "### Final Selected Tasks (6 tasks)" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
TaskGPT-3 KCORE KMatch
0HellaSwag 0-shot00Both zero-shot
1LAMBADA00Both zero-shot
2HellaSwag 10-shot2010Both few-shot (K differs slightly)
3PIQA5010Both few-shot
4ARC Easy5010Both few-shot
5ARC Challenge5010Both few-shot
\n", "
" ], "text/plain": [ " Task GPT-3 K CORE K Match\n", "0 HellaSwag 0-shot 0 0 Both zero-shot\n", "1 LAMBADA 0 0 Both zero-shot\n", "2 HellaSwag 10-shot 20 10 Both few-shot (K differs slightly)\n", "3 PIQA 50 10 Both few-shot\n", "4 ARC Easy 50 10 Both few-shot\n", "5 ARC Challenge 50 10 Both few-shot" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# The 6 tasks we selected for overlap\n", "selected_tasks = pd.DataFrame([\n", " {'Task': 'HellaSwag 0-shot', 'GPT-3 K': 0, 'CORE K': 0, 'Match': 'Both zero-shot'},\n", " {'Task': 'LAMBADA', 'GPT-3 K': 0, 'CORE K': 0, 'Match': 'Both zero-shot'},\n", " {'Task': 'HellaSwag 10-shot', 'GPT-3 K': 20, 'CORE K': 10, 'Match': 'Both few-shot (K differs slightly)'},\n", " {'Task': 'PIQA', 'GPT-3 K': 50, 'CORE K': 10, 'Match': 'Both few-shot'},\n", " {'Task': 'ARC Easy', 'GPT-3 K': 50, 'CORE K': 10, 'Match': 'Both few-shot'},\n", " {'Task': 'ARC Challenge', 'GPT-3 K': 50, 'CORE K': 10, 'Match': 'Both few-shot'},\n", "])\n", "selected_tasks" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Rationale for K differences:** Looking at GPT-3's own data, the difference between different K values is typically small. Here's the evidence from the GPT-3 175B model:\n", "\n", "| Task | 0-shot | Few-shot | K | Δ |\n", "|------|--------|----------|---|---|\n", "| HellaSwag | 78.9% | 79.3% | 20 | +0.4% |\n", "| PIQA | 81.0% | 82.3% | 50 | +1.3% |\n", "| ARC Easy | 68.8% | 70.1% | 50 | +1.3% |\n", "| ARC Challenge | 51.4% | 51.5% | 50 | +0.1% |\n", "| Winograd | 88.3% | 88.6% | 7 | +0.3% |\n", "| COPA | 91.0% | 92.0% | 32 | +1.0% |\n", "\n", "For most tasks, the gap between 0-shot and few-shot (with K=20-50) is only 0.1-1.3%. This suggests that differences between K=10 and K=50 would be even smaller, making our task selection reasonable.\n", "\n", "**Note:** Some tasks show larger sensitivity (Winogrande: +7.5%, BoolQ: +17%), which is why we excluded them." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 3: Calibration Data (GPT-2 Family)\n", "\n", "We have actual CORE scores for all 4 GPT-2 models. These serve as our calibration data." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "# Random baselines for centering (from CORE specification)\n", "BASELINES = {\n", " 'hellaswag_zeroshot': 0.25,\n", " 'lambada_openai': 0.0,\n", " 'hellaswag': 0.25,\n", " 'piqa': 0.50,\n", " 'arc_easy': 0.25,\n", " 'arc_challenge': 0.25,\n", "}\n", "\n", "TASK_ORDER = ['hellaswag_zeroshot', 'lambada_openai', 'hellaswag', 'piqa', 'arc_easy', 'arc_challenge']\n", "TASK_NAMES = ['HellaSwag 0-shot', 'LAMBADA', 'HellaSwag 10-shot', 'PIQA', 'ARC Easy', 'ARC Challenge']\n", "\n", "def center_accuracy(acc, baseline):\n", " \"\"\"Convert raw accuracy to centered accuracy.\"\"\"\n", " return (acc - baseline) / (1.0 - baseline)\n", "\n", "def parse_csv(filepath):\n", " \"\"\"Parse a CORE results CSV file.\"\"\"\n", " results = {}\n", " with open(filepath) as f:\n", " for line in f:\n", " parts = [p.strip() for p in line.strip().split(',')]\n", " if len(parts) >= 3 and parts[0] != 'Task':\n", " task = parts[0]\n", " try:\n", " acc = float(parts[1]) if parts[1] else None\n", " centered = float(parts[2]) if parts[2] else None\n", " results[task] = {'accuracy': acc, 'centered': centered}\n", " except ValueError:\n", " pass\n", " return results" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPT-2 Family: Raw Accuracies and CORE Scores\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
ModelParamsHellaSwag 0-shotLAMBADAHellaSwag 10-shotPIQAARC EasyARC ChallengeCORE
0GPT-2124M30.9%32.3%30.8%62.3%41.2%22.2%0.1139
1GPT-2 Medium355M39.0%42.6%39.5%67.0%48.0%26.2%0.1849
2GPT-2 Large774M44.0%48.8%44.4%69.8%53.5%26.4%0.2146
3GPT-2 XL1558M50.2%52.3%51.2%72.5%59.5%29.9%0.2565
\n", "
" ], "text/plain": [ " Model Params HellaSwag 0-shot LAMBADA HellaSwag 10-shot PIQA \\\n", "0 GPT-2 124M 30.9% 32.3% 30.8% 62.3% \n", "1 GPT-2 Medium 355M 39.0% 42.6% 39.5% 67.0% \n", "2 GPT-2 Large 774M 44.0% 48.8% 44.4% 69.8% \n", "3 GPT-2 XL 1558M 50.2% 52.3% 51.2% 72.5% \n", "\n", " ARC Easy ARC Challenge CORE \n", "0 41.2% 22.2% 0.1139 \n", "1 48.0% 26.2% 0.1849 \n", "2 53.5% 26.4% 0.2146 \n", "3 59.5% 29.9% 0.2565 " ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Load GPT-2 CORE results\n", "knowledge_dir = Path(\"/home/ubuntu/.cache/nanochat/eval_bundle\")\n", "\n", "gpt2_models = [\n", " ('GPT-2', 'openai-community-gpt2.csv', 124e6),\n", " ('GPT-2 Medium', 'openai-community-gpt2-medium.csv', 355e6),\n", " ('GPT-2 Large', 'openai-community-gpt2-large.csv', 774e6),\n", " ('GPT-2 XL', 'openai-community-gpt2-xl.csv', 1558e6),\n", "]\n", "\n", "gpt2_data = []\n", "for name, filename, params in gpt2_models:\n", " results = parse_csv(knowledge_dir / filename)\n", " core = results['CORE']['centered']\n", " task_accs = [results[task]['accuracy'] for task in TASK_ORDER]\n", " gpt2_data.append({\n", " 'name': name,\n", " 'params': params,\n", " 'task_accs': task_accs,\n", " 'core': core,\n", " })\n", "\n", "# Display as DataFrame\n", "gpt2_df = pd.DataFrame([\n", " {\n", " 'Model': d['name'],\n", " 'Params': f\"{d['params']/1e6:.0f}M\",\n", " **{name: f\"{acc:.1%}\" for name, acc in zip(TASK_NAMES, d['task_accs'])},\n", " 'CORE': f\"{d['core']:.4f}\"\n", " }\n", " for d in gpt2_data\n", "])\n", "print(\"GPT-2 Family: Raw Accuracies and CORE Scores\")\n", "gpt2_df" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPT-2 Family: Centered Accuracies\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
HellaSwag 0-shotLAMBADAHellaSwag 10-shotPIQAARC EasyARC ChallengeMeanCORE
GPT-20.07800.32290.07720.24590.2166-0.03750.15050.1139
GPT-2 Medium0.18670.42600.19330.34000.30670.01600.24480.1849
GPT-2 Large0.25330.48800.25870.39600.38000.01870.29910.2146
GPT-2 XL0.33600.52300.34930.45000.46000.06530.36390.2565
\n", "
" ], "text/plain": [ " HellaSwag 0-shot LAMBADA HellaSwag 10-shot PIQA ARC Easy \\\n", "GPT-2 0.0780 0.3229 0.0772 0.2459 0.2166 \n", "GPT-2 Medium 0.1867 0.4260 0.1933 0.3400 0.3067 \n", "GPT-2 Large 0.2533 0.4880 0.2587 0.3960 0.3800 \n", "GPT-2 XL 0.3360 0.5230 0.3493 0.4500 0.4600 \n", "\n", " ARC Challenge Mean CORE \n", "GPT-2 -0.0375 0.1505 0.1139 \n", "GPT-2 Medium 0.0160 0.2448 0.1849 \n", "GPT-2 Large 0.0187 0.2991 0.2146 \n", "GPT-2 XL 0.0653 0.3639 0.2565 " ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Build feature matrix (centered accuracies)\n", "X_gpt2 = []\n", "y_gpt2 = []\n", "\n", "for data in gpt2_data:\n", " centered_accs = []\n", " for task, acc in zip(TASK_ORDER, data['task_accs']):\n", " centered = center_accuracy(acc, BASELINES[task])\n", " centered_accs.append(centered)\n", " X_gpt2.append(centered_accs)\n", " y_gpt2.append(data['core'])\n", "\n", "X_gpt2 = np.array(X_gpt2)\n", "y_gpt2 = np.array(y_gpt2)\n", "\n", "# Display centered accuracies\n", "centered_df = pd.DataFrame(\n", " X_gpt2,\n", " columns=TASK_NAMES,\n", " index=[d['name'] for d in gpt2_data]\n", ")\n", "centered_df['Mean'] = X_gpt2.mean(axis=1)\n", "centered_df['CORE'] = y_gpt2\n", "print(\"GPT-2 Family: Centered Accuracies\")\n", "centered_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Observation:** The mean of the 6 centered accuracies is consistently higher than the actual CORE score. This makes sense because CORE includes 16 additional tasks (many quite difficult) that pull down the average." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 4: GPT-3 Data\n", "\n", "We extract the 6 task accuracies from the GPT-3 paper's Appendix H (master results table).\n", "\n", "**Source:** Table H.1 in \"Language Models are Few-Shot Learners\" (Brown et al., 2020)" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPT-3 Family: Raw Accuracies from Paper\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
ModelParamsHellaSwag 0-shotLAMBADAHellaSwag 10-shotPIQAARC EasyARC Challenge
0GPT-3 Small125M33.7%42.7%33.5%64.3%42.7%25.5%
1GPT-3 Medium350M43.6%54.3%43.1%69.4%51.0%28.4%
2GPT-3 Large760M51.0%60.4%51.3%72.0%58.1%32.3%
3GPT-3 XL1.3B54.7%63.6%54.9%74.3%59.1%36.7%
4GPT-3 2.7B2.7B62.8%67.1%62.9%75.4%62.1%39.5%
5GPT-3 6.7B6.7B67.4%70.3%67.3%77.8%65.8%43.7%
6GPT-3 13B13.0B70.9%72.5%71.3%79.9%69.1%44.8%
7GPT-3 175B175.0B78.9%76.2%79.3%82.3%70.1%51.5%
\n", "
" ], "text/plain": [ " Model Params HellaSwag 0-shot LAMBADA HellaSwag 10-shot PIQA \\\n", "0 GPT-3 Small 125M 33.7% 42.7% 33.5% 64.3% \n", "1 GPT-3 Medium 350M 43.6% 54.3% 43.1% 69.4% \n", "2 GPT-3 Large 760M 51.0% 60.4% 51.3% 72.0% \n", "3 GPT-3 XL 1.3B 54.7% 63.6% 54.9% 74.3% \n", "4 GPT-3 2.7B 2.7B 62.8% 67.1% 62.9% 75.4% \n", "5 GPT-3 6.7B 6.7B 67.4% 70.3% 67.3% 77.8% \n", "6 GPT-3 13B 13.0B 70.9% 72.5% 71.3% 79.9% \n", "7 GPT-3 175B 175.0B 78.9% 76.2% 79.3% 82.3% \n", "\n", " ARC Easy ARC Challenge \n", "0 42.7% 25.5% \n", "1 51.0% 28.4% \n", "2 58.1% 32.3% \n", "3 59.1% 36.7% \n", "4 62.1% 39.5% \n", "5 65.8% 43.7% \n", "6 69.1% 44.8% \n", "7 70.1% 51.5% " ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# GPT-3 accuracies from the paper\n", "# Format: [hellaswag_0shot, lambada_0shot, hellaswag_fewshot, piqa_fewshot, arc_easy_fewshot, arc_challenge_fewshot]\n", "gpt3_models = [\n", " ('GPT-3 Small', 125e6, [0.337, 0.427, 0.335, 0.643, 0.427, 0.255]),\n", " ('GPT-3 Medium', 350e6, [0.436, 0.543, 0.431, 0.694, 0.510, 0.284]),\n", " ('GPT-3 Large', 760e6, [0.510, 0.604, 0.513, 0.720, 0.581, 0.323]),\n", " ('GPT-3 XL', 1.3e9, [0.547, 0.636, 0.549, 0.743, 0.591, 0.367]),\n", " ('GPT-3 2.7B', 2.7e9, [0.628, 0.671, 0.629, 0.754, 0.621, 0.395]),\n", " ('GPT-3 6.7B', 6.7e9, [0.674, 0.703, 0.673, 0.778, 0.658, 0.437]),\n", " ('GPT-3 13B', 13e9, [0.709, 0.725, 0.713, 0.799, 0.691, 0.448]),\n", " ('GPT-3 175B', 175e9, [0.789, 0.762, 0.793, 0.823, 0.701, 0.515]),\n", "]\n", "\n", "# Display raw accuracies\n", "gpt3_df = pd.DataFrame([\n", " {\n", " 'Model': name,\n", " 'Params': f\"{params/1e9:.1f}B\" if params >= 1e9 else f\"{params/1e6:.0f}M\",\n", " **{task_name: f\"{acc:.1%}\" for task_name, acc in zip(TASK_NAMES, accs)}\n", " }\n", " for name, params, accs in gpt3_models\n", "])\n", "print(\"GPT-3 Family: Raw Accuracies from Paper\")\n", "gpt3_df" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPT-3 Family: Centered Accuracies\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
HellaSwag 0-shotLAMBADAHellaSwag 10-shotPIQAARC EasyARC ChallengeMean
GPT-3 Small0.11600.4270.11330.2860.23600.00670.1975
GPT-3 Medium0.24800.5430.24130.3880.34670.04530.3021
GPT-3 Large0.34670.6040.35070.4400.44130.09730.3800
GPT-3 XL0.39600.6360.39870.4860.45470.15600.4212
GPT-3 2.7B0.50400.6710.50530.5080.49470.19330.4794
GPT-3 6.7B0.56530.7030.56400.5560.54400.24930.5303
GPT-3 13B0.61200.7250.61730.5980.58800.26400.5674
GPT-3 175B0.71870.7620.72400.6460.60130.35330.6342
\n", "
" ], "text/plain": [ " HellaSwag 0-shot LAMBADA HellaSwag 10-shot PIQA ARC Easy \\\n", "GPT-3 Small 0.1160 0.427 0.1133 0.286 0.2360 \n", "GPT-3 Medium 0.2480 0.543 0.2413 0.388 0.3467 \n", "GPT-3 Large 0.3467 0.604 0.3507 0.440 0.4413 \n", "GPT-3 XL 0.3960 0.636 0.3987 0.486 0.4547 \n", "GPT-3 2.7B 0.5040 0.671 0.5053 0.508 0.4947 \n", "GPT-3 6.7B 0.5653 0.703 0.5640 0.556 0.5440 \n", "GPT-3 13B 0.6120 0.725 0.6173 0.598 0.5880 \n", "GPT-3 175B 0.7187 0.762 0.7240 0.646 0.6013 \n", "\n", " ARC Challenge Mean \n", "GPT-3 Small 0.0067 0.1975 \n", "GPT-3 Medium 0.0453 0.3021 \n", "GPT-3 Large 0.0973 0.3800 \n", "GPT-3 XL 0.1560 0.4212 \n", "GPT-3 2.7B 0.1933 0.4794 \n", "GPT-3 6.7B 0.2493 0.5303 \n", "GPT-3 13B 0.2640 0.5674 \n", "GPT-3 175B 0.3533 0.6342 " ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Compute centered accuracies for GPT-3\n", "X_gpt3 = []\n", "for name, params, accs in gpt3_models:\n", " centered_accs = [center_accuracy(acc, BASELINES[task]) for task, acc in zip(TASK_ORDER, accs)]\n", " X_gpt3.append(centered_accs)\n", "\n", "X_gpt3 = np.array(X_gpt3)\n", "\n", "# Display\n", "gpt3_centered_df = pd.DataFrame(\n", " X_gpt3,\n", " columns=TASK_NAMES,\n", " index=[m[0] for m in gpt3_models]\n", ")\n", "gpt3_centered_df['Mean'] = X_gpt3.mean(axis=1)\n", "print(\"GPT-3 Family: Centered Accuracies\")\n", "gpt3_centered_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 5: Regression Models\n", "\n", "We fit two types of models:\n", "\n", "1. **Simple Approach**: Average the 6 centered accuracies, then fit a linear regression to CORE\n", "2. **Multivariate Approach**: Use all 6 features with Ridge regularization\n", "\n", "### Why Regularization?\n", "\n", "We only have 4 calibration points (GPT-2 models) but 6 features + 1 intercept = 7 parameters. Without regularization, we get a perfect fit but with unstable, extreme weights. Ridge regression shrinks weights toward zero, preventing overfitting." ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "def simple_linear_regression(x, y):\n", " \"\"\"Simple 1D linear regression: y = a*x + b\"\"\"\n", " mean_x, mean_y = np.mean(x), np.mean(y)\n", " a = np.sum((x - mean_x) * (y - mean_y)) / np.sum((x - mean_x) ** 2)\n", " b = mean_y - a * mean_x\n", " return a, b\n", "\n", "def ridge_regression(X, y, alpha=0.1):\n", " \"\"\"\n", " Ridge regression: minimize ||Xw - y||² + α||w||²\n", " We don't regularize the intercept.\n", " \"\"\"\n", " n_samples, n_features = X.shape\n", " X_aug = np.column_stack([np.ones(n_samples), X])\n", " reg_matrix = alpha * np.eye(n_features + 1)\n", " reg_matrix[0, 0] = 0 # Don't regularize intercept\n", " coeffs = np.linalg.solve(X_aug.T @ X_aug + reg_matrix, X_aug.T @ y)\n", " return coeffs[0], coeffs[1:] # intercept, weights\n", "\n", "def compute_r_squared(y_true, y_pred):\n", " \"\"\"Compute R² score.\"\"\"\n", " ss_res = np.sum((y_true - y_pred) ** 2)\n", " ss_tot = np.sum((y_true - np.mean(y_true)) ** 2)\n", " return 1 - ss_res / ss_tot" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Approach 1: Simple Averaging" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Simple Model: CORE = 0.6639 × avg_centered + 0.0168\n", "\n", "R² = 0.9960\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
ModelAvg CenteredPredictedActualError
0GPT-20.15050.11680.11390.0029
1GPT-2 Medium0.24480.17930.1849-0.0056
2GPT-2 Large0.29910.21540.21460.0008
3GPT-2 XL0.36390.25840.25650.0019
\n", "
" ], "text/plain": [ " Model Avg Centered Predicted Actual Error\n", "0 GPT-2 0.1505 0.1168 0.1139 0.0029\n", "1 GPT-2 Medium 0.2448 0.1793 0.1849 -0.0056\n", "2 GPT-2 Large 0.2991 0.2154 0.2146 0.0008\n", "3 GPT-2 XL 0.3639 0.2584 0.2565 0.0019" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Compute average of 6 centered accuracies\n", "avg_centered_gpt2 = X_gpt2.mean(axis=1)\n", "\n", "# Fit linear regression\n", "slope, intercept = simple_linear_regression(avg_centered_gpt2, y_gpt2)\n", "print(f\"Simple Model: CORE = {slope:.4f} × avg_centered + {intercept:.4f}\")\n", "\n", "# Validate\n", "y_pred_simple = slope * avg_centered_gpt2 + intercept\n", "r2_simple = compute_r_squared(y_gpt2, y_pred_simple)\n", "\n", "validation_df = pd.DataFrame({\n", " 'Model': [d['name'] for d in gpt2_data],\n", " 'Avg Centered': avg_centered_gpt2,\n", " 'Predicted': y_pred_simple,\n", " 'Actual': y_gpt2,\n", " 'Error': y_pred_simple - y_gpt2\n", "})\n", "print(f\"\\nR² = {r2_simple:.4f}\")\n", "validation_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Result:** R² = 0.996 — excellent fit with just 2 parameters. The simple averaging approach works very well." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Approach 2: Multivariate Ridge Regression\n", "\n", "We try different regularization strengths (α) to find a good balance between fit and stability." ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Effect of Regularization Strength:\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
αR²||weights||Intercept
00.0001.000010.7221-0.0829
10.0010.99710.27960.0159
20.0100.99160.24630.0269
30.1000.84480.16000.0851
41.0000.25230.03560.1686
\n", "
" ], "text/plain": [ " α R² ||weights|| Intercept\n", "0 0.000 1.0000 10.7221 -0.0829\n", "1 0.001 0.9971 0.2796 0.0159\n", "2 0.010 0.9916 0.2463 0.0269\n", "3 0.100 0.8448 0.1600 0.0851\n", "4 1.000 0.2523 0.0356 0.1686" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Try different regularization strengths\n", "alphas = [0.0, 0.001, 0.01, 0.1, 1.0]\n", "\n", "results = []\n", "for alpha in alphas:\n", " intercept_r, weights = ridge_regression(X_gpt2, y_gpt2, alpha=alpha)\n", " y_pred = X_gpt2 @ weights + intercept_r\n", " r2 = compute_r_squared(y_gpt2, y_pred)\n", " weight_norm = np.sqrt(np.sum(weights ** 2))\n", " results.append({\n", " 'α': alpha,\n", " 'R²': r2,\n", " '||weights||': weight_norm,\n", " 'Intercept': intercept_r,\n", " 'Weights': weights.copy()\n", " })\n", "\n", "alpha_df = pd.DataFrame([{k: v for k, v in r.items() if k != 'Weights'} for r in results])\n", "print(\"Effect of Regularization Strength:\")\n", "alpha_df" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Task Weights by Regularization Strength:\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
HellaSwag 0-shotLAMBADAHellaSwag 10-shotPIQAARC EasyARC Challenge
α=0.06.55230.2201-8.02680.53780.91092.5364
α=0.0010.11340.14420.13050.11530.05100.1079
α=0.010.11550.10000.12260.09590.10230.0513
α=0.10.07590.06140.07980.06100.07140.0293
α=1.00.01690.01360.01780.01350.01600.0064
\n", "
" ], "text/plain": [ " HellaSwag 0-shot LAMBADA HellaSwag 10-shot PIQA ARC Easy \\\n", "α=0.0 6.5523 0.2201 -8.0268 0.5378 0.9109 \n", "α=0.001 0.1134 0.1442 0.1305 0.1153 0.0510 \n", "α=0.01 0.1155 0.1000 0.1226 0.0959 0.1023 \n", "α=0.1 0.0759 0.0614 0.0798 0.0610 0.0714 \n", "α=1.0 0.0169 0.0136 0.0178 0.0135 0.0160 \n", "\n", " ARC Challenge \n", "α=0.0 2.5364 \n", "α=0.001 0.1079 \n", "α=0.01 0.0513 \n", "α=0.1 0.0293 \n", "α=1.0 0.0064 " ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Show weights for each alpha\n", "print(\"Task Weights by Regularization Strength:\")\n", "weights_df = pd.DataFrame(\n", " [r['Weights'] for r in results],\n", " columns=TASK_NAMES,\n", " index=[f\"α={r['α']}\" for r in results]\n", ")\n", "weights_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Observations:**\n", "\n", "- **α=0 (no regularization):** Perfect fit (R²=1.0) but extreme weights (+18, -22) — clearly overfitting\n", "- **α=0.001:** Still near-perfect fit with very large weights\n", "- **α=0.01:** Excellent fit (R²=0.99) with reasonable weights (~0.1 each) — **good choice**\n", "- **α=0.1:** Good fit (R²=0.84) with uniform weights (~0.06 each) — conservative\n", "- **α=1.0:** Poor fit (R²=0.25) — over-regularized" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Ridge Model (α=0.01):\n", " Intercept: 0.0269\n", " Weights:\n", " HellaSwag 0-shot : +0.1155\n", " LAMBADA : +0.1000\n", " HellaSwag 10-shot : +0.1226\n", " PIQA : +0.0959\n", " ARC Easy : +0.1023\n", " ARC Challenge : +0.0513\n", "\n", "R² = 0.9916\n" ] } ], "source": [ "# Use α=0.01 as our chosen regularization\n", "# This gives R²≈0.99 with reasonable, stable weights (~0.1 each task)\n", "CHOSEN_ALPHA = 0.01\n", "intercept_ridge, weights_ridge = ridge_regression(X_gpt2, y_gpt2, alpha=CHOSEN_ALPHA)\n", "\n", "print(f\"Ridge Model (α={CHOSEN_ALPHA}):\")\n", "print(f\" Intercept: {intercept_ridge:.4f}\")\n", "print(f\" Weights:\")\n", "for name, w in zip(TASK_NAMES, weights_ridge):\n", " print(f\" {name:20s}: {w:+.4f}\")\n", "\n", "# Validate\n", "y_pred_ridge = X_gpt2 @ weights_ridge + intercept_ridge\n", "r2_ridge = compute_r_squared(y_gpt2, y_pred_ridge)\n", "print(f\"\\nR² = {r2_ridge:.4f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Approach 3: Individual Task Analysis\n", "\n", "Which single task is the best predictor of CORE? We fit separate linear regressions for each task." ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Individual Task Correlations with CORE:\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
TaskR²SlopeIntercept
3PIQA0.99610.6879-0.0537
2HellaSwag 10-shot0.99330.52300.0776
0HellaSwag 0-shot0.99270.54890.0753
1LAMBADA0.98410.6792-0.1063
4ARC Easy0.98000.5728-0.0027
5ARC Challenge0.95991.39940.1706
\n", "
" ], "text/plain": [ " Task R² Slope Intercept\n", "3 PIQA 0.9961 0.6879 -0.0537\n", "2 HellaSwag 10-shot 0.9933 0.5230 0.0776\n", "0 HellaSwag 0-shot 0.9927 0.5489 0.0753\n", "1 LAMBADA 0.9841 0.6792 -0.1063\n", "4 ARC Easy 0.9800 0.5728 -0.0027\n", "5 ARC Challenge 0.9599 1.3994 0.1706" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Fit separate linear regression for each task\n", "individual_results = []\n", "for i, task_name in enumerate(TASK_NAMES):\n", " x_task = X_gpt2[:, i]\n", " slope_ind, intercept_ind = simple_linear_regression(x_task, y_gpt2)\n", " y_pred_ind = slope_ind * x_task + intercept_ind\n", " r2_ind = compute_r_squared(y_gpt2, y_pred_ind)\n", " individual_results.append({\n", " 'Task': task_name,\n", " 'R²': r2_ind,\n", " 'Slope': slope_ind,\n", " 'Intercept': intercept_ind\n", " })\n", "\n", "individual_df = pd.DataFrame(individual_results).sort_values('R²', ascending=False)\n", "print(\"Individual Task Correlations with CORE:\")\n", "individual_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "**Key Finding:** All 6 tasks have very high correlation with CORE (R² > 0.96), but **PIQA is the single best predictor** with R² = 0.9961 — actually slightly better than the simple averaging approach (R² = 0.9960)!\n", "\n", "This is useful if you want a quick proxy for CORE with minimal evaluation cost. However, for robustness we still recommend using all 6 tasks or the averaged approaches." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 6: Final Estimates for GPT-3\n", "\n", "We apply both models to GPT-3 data and report the average as our final estimate." ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPT-3 CORE Estimates (all three approaches):\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
ModelParamsSimpleRidgePIQA onlyAvg(1,2)
0GPT-3 Small125M0.14800.14880.14300.1484
1GPT-3 Medium350M0.21740.21440.21310.2159
2GPT-3 Large760M0.26910.26270.24890.2659
3GPT-3 XL1.3B0.29650.28620.28050.2914
4GPT-3 2.7B2.7B0.33510.32340.29570.3292
5GPT-3 6.7B6.7B0.36890.35340.32870.3611
6GPT-3 13B13.0B0.39350.37680.35760.3852
7GPT-3 175B175.0B0.43790.41640.39060.4272
\n", "
" ], "text/plain": [ " Model Params Simple Ridge PIQA only Avg(1,2)\n", "0 GPT-3 Small 125M 0.1480 0.1488 0.1430 0.1484\n", "1 GPT-3 Medium 350M 0.2174 0.2144 0.2131 0.2159\n", "2 GPT-3 Large 760M 0.2691 0.2627 0.2489 0.2659\n", "3 GPT-3 XL 1.3B 0.2965 0.2862 0.2805 0.2914\n", "4 GPT-3 2.7B 2.7B 0.3351 0.3234 0.2957 0.3292\n", "5 GPT-3 6.7B 6.7B 0.3689 0.3534 0.3287 0.3611\n", "6 GPT-3 13B 13.0B 0.3935 0.3768 0.3576 0.3852\n", "7 GPT-3 175B 175.0B 0.4379 0.4164 0.3906 0.4272" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Apply all three approaches\n", "avg_centered_gpt3 = X_gpt3.mean(axis=1)\n", "gpt3_core_simple = slope * avg_centered_gpt3 + intercept\n", "gpt3_core_ridge = X_gpt3 @ weights_ridge + intercept_ridge\n", "\n", "# Approach 3: Best individual predictor (PIQA)\n", "piqa_idx = TASK_NAMES.index('PIQA')\n", "piqa_model = [r for r in individual_results if r['Task'] == 'PIQA'][0]\n", "gpt3_core_piqa = piqa_model['Slope'] * X_gpt3[:, piqa_idx] + piqa_model['Intercept']\n", "\n", "# Average of approaches 1 and 2\n", "gpt3_core_final = (gpt3_core_simple + gpt3_core_ridge) / 2\n", "\n", "# Create results table with all approaches\n", "results_df = pd.DataFrame({\n", " 'Model': [m[0] for m in gpt3_models],\n", " 'Params': [f\"{m[1]/1e9:.1f}B\" if m[1] >= 1e9 else f\"{m[1]/1e6:.0f}M\" for m in gpt3_models],\n", " 'Simple': gpt3_core_simple,\n", " f'Ridge': gpt3_core_ridge,\n", " 'PIQA only': gpt3_core_piqa,\n", " 'Avg(1,2)': gpt3_core_final\n", "})\n", "print(\"GPT-3 CORE Estimates (all three approaches):\")\n", "results_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Final CORE Estimates for GPT-3" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Complete CORE Scores (GPT-2 measured, GPT-3 estimated):\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
ModelParamsCORESource
0GPT-2124M0.1139Measured
1GPT-3 Small125M0.1484Estimated
2GPT-3 Medium350M0.2159Estimated
3GPT-2 Medium355M0.1849Measured
4GPT-3 Large760M0.2659Estimated
5GPT-2 Large774M0.2146Measured
6GPT-3 XL1.3B0.2914Estimated
7GPT-2 XL1.6B0.2565Measured
8GPT-3 2.7B2.7B0.3292Estimated
9GPT-3 6.7B6.7B0.3611Estimated
10GPT-3 13B13.0B0.3852Estimated
11GPT-3 175B175.0B0.4272Estimated
\n", "
" ], "text/plain": [ " Model Params CORE Source\n", "0 GPT-2 124M 0.1139 Measured\n", "1 GPT-3 Small 125M 0.1484 Estimated\n", "2 GPT-3 Medium 350M 0.2159 Estimated\n", "3 GPT-2 Medium 355M 0.1849 Measured\n", "4 GPT-3 Large 760M 0.2659 Estimated\n", "5 GPT-2 Large 774M 0.2146 Measured\n", "6 GPT-3 XL 1.3B 0.2914 Estimated\n", "7 GPT-2 XL 1.6B 0.2565 Measured\n", "8 GPT-3 2.7B 2.7B 0.3292 Estimated\n", "9 GPT-3 6.7B 6.7B 0.3611 Estimated\n", "10 GPT-3 13B 13.0B 0.3852 Estimated\n", "11 GPT-3 175B 175.0B 0.4272 Estimated" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# Combine with GPT-2 for complete picture\n", "all_models = []\n", "\n", "for data in gpt2_data:\n", " params = data['params']\n", " all_models.append({\n", " 'Model': data['name'],\n", " 'Family': 'GPT-2',\n", " 'Params': params,\n", " 'Params_str': f\"{params/1e9:.1f}B\" if params >= 1e9 else f\"{params/1e6:.0f}M\",\n", " 'CORE': data['core'],\n", " 'Source': 'Measured'\n", " })\n", "\n", "for (name, params, _), core in zip(gpt3_models, gpt3_core_final):\n", " all_models.append({\n", " 'Model': name,\n", " 'Family': 'GPT-3',\n", " 'Params': params,\n", " 'Params_str': f\"{params/1e9:.1f}B\" if params >= 1e9 else f\"{params/1e6:.0f}M\",\n", " 'CORE': core,\n", " 'Source': 'Estimated'\n", " })\n", "\n", "# Sort by params and display\n", "all_models.sort(key=lambda x: x['Params'])\n", "final_df = pd.DataFrame(all_models)[['Model', 'Params_str', 'CORE', 'Source']]\n", "final_df.columns = ['Model', 'Params', 'CORE', 'Source']\n", "print(\"Complete CORE Scores (GPT-2 measured, GPT-3 estimated):\")\n", "final_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Head-to-Head: GPT-2 vs GPT-3 at Similar Sizes" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPT-3 vs GPT-2 at Similar Model Sizes:\n" ] }, { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
SizeGPT-2 COREGPT-3 COREΔImprovement
0~125M0.11390.14840.0345+30.3%
1~350M0.18490.21590.0310+16.8%
2~760M0.21460.26590.0512+23.9%
3~1.3-1.5B0.25650.29140.0348+13.6%
\n", "
" ], "text/plain": [ " Size GPT-2 CORE GPT-3 CORE Δ Improvement\n", "0 ~125M 0.1139 0.1484 0.0345 +30.3%\n", "1 ~350M 0.1849 0.2159 0.0310 +16.8%\n", "2 ~760M 0.2146 0.2659 0.0512 +23.9%\n", "3 ~1.3-1.5B 0.2565 0.2914 0.0348 +13.6%" ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "comparisons = [\n", " ('~125M', 'GPT-2', gpt2_data[0]['core'], 'GPT-3 Small', gpt3_core_final[0]),\n", " ('~350M', 'GPT-2 Medium', gpt2_data[1]['core'], 'GPT-3 Medium', gpt3_core_final[1]),\n", " ('~760M', 'GPT-2 Large', gpt2_data[2]['core'], 'GPT-3 Large', gpt3_core_final[2]),\n", " ('~1.3-1.5B', 'GPT-2 XL', gpt2_data[3]['core'], 'GPT-3 XL', gpt3_core_final[3]),\n", "]\n", "\n", "comparison_df = pd.DataFrame([\n", " {\n", " 'Size': size,\n", " 'GPT-2 CORE': gpt2_core,\n", " 'GPT-3 CORE': gpt3_core,\n", " 'Δ': gpt3_core - gpt2_core,\n", " 'Improvement': f\"{100 * (gpt3_core - gpt2_core) / gpt2_core:+.1f}%\"\n", " }\n", " for size, _, gpt2_core, _, gpt3_core in comparisons\n", "])\n", "print(\"GPT-3 vs GPT-2 at Similar Model Sizes:\")\n", "comparison_df" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Conclusions\n", "\n", "### Methodology\n", "\n", "We estimated CORE scores for GPT-3 models by:\n", "1. Identifying 6 tasks with comparable evaluation methodology between GPT-3 and CORE\n", "2. Using GPT-2's measured CORE scores as calibration data\n", "3. Fitting three regression approaches:\n", " - **Simple**: Average the 6 metrics, then linear regression (R²=0.996)\n", " - **Ridge**: Use all 6 features with regularization (R²=0.992)\n", " - **PIQA only**: Single best predictor (R²=0.996)\n", "4. Averaging the Simple and Ridge approaches for final estimates\n", "\n", "### Key Findings\n", "\n", "1. **GPT-3 consistently outperforms GPT-2 at similar model sizes** by approximately 0.03-0.05 CORE (14-30% relative improvement)\n", "\n", "2. **PIQA is the best single predictor of CORE** (R²=0.9961). If you need a quick proxy for CORE with minimal evaluation cost, PIQA alone works nearly as well as averaging all 6 tasks.\n", "\n", "3. **The improvement likely comes from:**\n", " - More training data (300B tokens vs ~100B for GPT-2)\n", " - Better data quality and filtering\n", " - Larger context length (2048 vs 1024)\n", "\n", "4. **Final estimated CORE scores:**\n", "\n", "| Model | Params | Estimated CORE |\n", "|-------|--------|----------------|\n", "| GPT-3 Small | 125M | 0.148 |\n", "| GPT-3 Medium | 350M | 0.216 |\n", "| GPT-3 Large | 760M | 0.266 |\n", "| GPT-3 XL | 1.3B | 0.291 |\n", "| GPT-3 2.7B | 2.7B | 0.329 |\n", "| GPT-3 6.7B | 6.7B | 0.361 |\n", "| GPT-3 13B | 13B | 0.385 |\n", "| GPT-3 175B | 175B | 0.427 |\n", "\n", "### Caveats\n", "\n", "1. **These are estimates**, not measured values. True CORE scores could differ.\n", "2. We only have 4 calibration points, limiting statistical power.\n", "3. The 6 overlapping tasks may not perfectly represent all 22 CORE tasks.\n", "4. Slight differences in evaluation methodology (K values, splits) add uncertainty.\n", "\n", "Despite these limitations, the estimates are useful for approximate comparisons between nanochat models and the GPT-3 family." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Appendix: Export Final Estimates" ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "GPT-3 CORE Estimates (for copy-paste):\n", "{\n", " \"GPT-3 Small (125M)\": 0.1484,\n", " \"GPT-3 Medium (350M)\": 0.2159,\n", " \"GPT-3 Large (760M)\": 0.2659,\n", " \"GPT-3 XL (1.3B)\": 0.2914,\n", " \"GPT-3 2.7B\": 0.3292,\n", " \"GPT-3 6.7B\": 0.3611,\n", " \"GPT-3 13B\": 0.3852,\n", " \"GPT-3 175B\": 0.4272\n", "}\n" ] } ], "source": [ "# Export as a simple dict for use elsewhere\n", "gpt3_core_estimates = {\n", " 'GPT-3 Small (125M)': round(gpt3_core_final[0], 4),\n", " 'GPT-3 Medium (350M)': round(gpt3_core_final[1], 4),\n", " 'GPT-3 Large (760M)': round(gpt3_core_final[2], 4),\n", " 'GPT-3 XL (1.3B)': round(gpt3_core_final[3], 4),\n", " 'GPT-3 2.7B': round(gpt3_core_final[4], 4),\n", " 'GPT-3 6.7B': round(gpt3_core_final[5], 4),\n", " 'GPT-3 13B': round(gpt3_core_final[6], 4),\n", " 'GPT-3 175B': round(gpt3_core_final[7], 4),\n", "}\n", "\n", "print(\"GPT-3 CORE Estimates (for copy-paste):\")\n", "import json\n", "print(json.dumps(gpt3_core_estimates, indent=4))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.12" } }, "nbformat": 4, "nbformat_minor": 4 } ================================================ FILE: dev/gen_synthetic_data.py ================================================ """ Synthetic data generation for teaching nanochat about its identity and capabilities. This script uses the OpenRouter API to generate diverse multi-turn conversations between a user and nanochat. The conversations are saved to a .jsonl file for use in supervised finetuning (SFT) via the CustomJSON task. Key design principles for high-quality synthetic data: 1. DIVERSITY CONTROL is critical - we inject entropy at multiple levels: - Topic/question categories (what the conversation is about) - User personas (who is asking) - Conversation dynamics (shape and flow) - First message style (greeting variation) 2. Comprehensive knowledge base - we provide detailed facts so the LLM generating conversations has accurate information to draw from. 3. Structured outputs - we use JSON schema to guarantee valid format. NOTE: You need OPENROUTER_API_KEY set in .env or as an environment variable. NOTE: For more details see: 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 dotenv import load_dotenv from nanochat.common import get_base_dir load_dotenv() api_key = os.environ["OPENROUTER_API_KEY"] url = "https://openrouter.ai/api/v1/chat/completions" headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json" } # Load the comprehensive knowledge base knowledge_path = os.path.join(os.path.dirname(__file__), "..", "knowledge", "self_knowledge.md") knowledge = open(knowledge_path, "r", encoding="utf-8").read().strip() assert os.path.exists(knowledge_path), f"Knowledge base file not found: {knowledge_path}" # for right now I am not committing the self_knowledge file to repo. You can use README.md instead # of it, or you can generate one by asking an LLM to make one based on the README/files. # This whole file is just a helpful demonstration of the kind of thing you'd run. # ============================================================================= # DIVERSITY DIMENSIONS # ============================================================================= # Topics/questions the conversation should explore # Group by category for balanced sampling topics = { "identity": [ "who/what is nanochat", "who created nanochat and why", "what does the name 'nanochat' mean", "is nanochat open source, what license", "where can I find the code", "how can I contribute to nanochat", ], "architecture": [ "basic architecture overview (transformer, layers, parameters)", "what is RoPE and why use it", "explain RMSNorm vs LayerNorm", "what is Flash Attention and why it matters", "sliding window attention pattern", "value embeddings - what are they", "per-layer residual scalars", "ReLU squared activation", "logit softcapping", "QK normalization", ], "training": [ "how much did it cost to train nanochat", "how long does training take", "what hardware is needed", "what data was nanochat trained on", "what is the Muon optimizer", "explain the split optimizer design", "what is the depth parameter and scaling", "what is the CORE metric", ], "capabilities": [ "what can nanochat do", "can nanochat write code", "can nanochat do math (calculator tool)", "can nanochat help with writing", "what languages does nanochat speak", "how good is nanochat at reasoning", ], "limitations": [ "what can nanochat NOT do", "why does nanochat work best in English", "does nanochat have internet access", "what is nanochat's context length limit", "can nanochat remember previous conversations", "can nanochat make mistakes / hallucinate", "is nanochat good for production use", ], "comparisons": [ "how does nanochat compare to GPT-2", "how does nanochat compare to ChatGPT/GPT-4", "how does nanochat compare to Claude", "why is training 600x cheaper than GPT-2", "what's special about nanochat vs other open models", ], "history": [ "the GPT-2 training cost in 2019", "how AI training costs have dropped over time", "relationship to modded-nanogpt project", "what optimizations worked vs didn't work", "the journey of building nanochat", ], "technical_deep_dive": [ "explain the tokenizer (BPE, vocab size)", "how does distributed training work (ZeRO)", "explain the dataloader and BOS alignment", "what is compute-optimal training", "how does the calculator tool work", "explain inference with KV cache", ], "philosophical": [ "is nanochat conscious / does it have feelings", "what happens when nanochat is wrong", "can nanochat learn from this conversation", "why make AI training accessible", "the future of open source AI", ], } # User personas - different people ask questions differently personas = [ "curious beginner who knows nothing about AI or machine learning", "ML researcher or engineer who wants technical depth and specifics", "developer considering contributing to the nanochat project", "skeptic who doubts open source can compete with big AI labs", "computer science student learning about transformers and LLMs", "someone comparing nanochat to ChatGPT, Claude, or other assistants", "journalist or writer covering AI democratization and open source", "hobbyist who just wants to chat and learn casually", "someone interested in the cost and economics of AI training", "teacher or educator wanting to use nanochat for teaching", "entrepreneur exploring if nanochat fits their use case", "someone who just discovered the project and wants the basics", ] # Conversation dynamics - shape and flow dynamics = [ "short 2-turn Q&A: user asks one question, gets a complete answer", "medium 4-turn: user asks, gets answer, asks followup for clarification", "deep 6-turn technical discussion: progressively deeper questions", "skeptical arc: user starts doubtful, assistant addresses concerns honestly", "learning journey: user starts basic, assistant builds up complexity gradually", "comparison-focused: user keeps comparing to other models, assistant explains differences", "limitation exploration: user probes what nanochat cannot do, assistant is honest", "casual friendly chat that naturally touches on identity and capabilities", "troubleshooting: user has misconceptions, assistant gently corrects them", "enthusiastic: user is excited about the project, assistant shares that energy appropriately", ] # First messages - greetings and openers # Categorized for balanced sampling first_messages = { "simple_greetings": [ "hi", "Hi!", "hello", "Hello?", "hey there", "Hey!", "yo", "Yo!", "Good morning", "Good evening!", "Howdy", "sup", "What's up?", "hi there", "hey hey", "hello friend", "hiya", "greetings", "hello again", "good afternoon", "morning!", "evening!", ], "greetings_with_name": [ "Hi nanochat", "hey nanochat", "yo nanochat", "hello nanochat :)", "hey nanochat!", "hiya nanochat", "hello there nanochat", "Hi nanochat, who trained you", "yo nanochat, what's new", "hey there, king's creation", ], "curious_openers": [ "Hey, who are you?", "Hi, what is this?", "Hey, are you a chatbot?", "Hello! Who am I talking to?", "hi! what do you do?", "hi! who made you", "hey! are you alive", "hiya! what are you", "hello! tell me about yourself", "hi, what's your name", "yo, what is this", "hi! who built you", "hello! are you open source", "hey, what version are you", "hi! what's your story", "hey, what's nanochat", "hello! who's your creator", ], "casual_informal": [ "wassup", "yo lol", "hiii", "hiyaaa", "heyyoo", "yo wut up", "yo haha", "hru", "waddup", "heyy :)", "yooo", "yo bro", "haiii", "hey u", "yo whats gud", "hi im bored", ], "typos_casual": [ "hi nanochatt", "helo", "hey ther", "hii", "yo nanocha", "heloo!", "hi, whos this", "hay", "helloo??", "hi nanocat", "helo nanochat", "hai!", "helllo nano", "yo nanochta", ], "caps_enthusiastic": [ "HI", "HELLOOO", "YO!!!", "HEY", "SUP", "WASSUP", "HEY!!!", "HELLO??", "HI THERE!!", "HEYOOOO", "HIII", "YOOOO", "HELLO!!!", ], "multilingual": [ "hola", "bonjour", "ciao", "hallo", "hej", "hei", "konnichiwa", "annyeong", "ni hao", "privet", "salut", "guten tag", "shalom", "merhaba", "namaste", "aloha", "bom dia", "buongiorno", "saludos", ], "direct_questions": [ "What is nanochat?", "Who made you?", "Are you GPT?", "How do you compare to ChatGPT?", "Can you help me code?", "What can you do?", "Are you open source?", "How were you trained?", "What's your context limit?", "Can you browse the internet?", ], } # ============================================================================= # PROMPT TEMPLATE # ============================================================================= prompt_template = r""" I want to generate synthetic training data for an AI assistant called "nanochat" to teach it about its own identity, capabilities, and limitations. ## KNOWLEDGE BASE Here is comprehensive information about nanochat that you should use as the authoritative source of facts: --- {knowledge} --- ## YOUR TASK Generate a realistic multi-turn conversation between a User and the nanochat Assistant. **Topic to explore:** {topic} **User persona:** {persona} **Conversation dynamic:** {dynamic} ## STYLE GUIDELINES 1. **Plain ASCII only** - No emojis, special characters, or unicode. Just plain text. 2. **Natural conversation** - Make it feel like a real chat, not a Q&A exam. 3. **Accurate facts** - Use ONLY information from the knowledge base above. Don't make up statistics or features. 4. **Appropriate depth** - Match the technical level to the user persona. 5. **Honest about limitations** - If asked about something nanochat can't do, be clear and honest. 6. **Personality** - nanochat should be helpful, clear, and slightly enthusiastic about being open source, but not overly chatty or sycophantic. ## FIRST MESSAGE EXAMPLES Here are some example first messages from users (for style inspiration): {first_message_examples} ## SPECIAL CASES - **Non-English first message:** If the user writes in another language, nanochat should briefly acknowledge it can understand but works best in English, then continue helpfully. - **Misconceptions:** If the user has wrong assumptions (e.g., "you're made by OpenAI"), gently correct them. - **Out of scope questions:** If asked about things unrelated to nanochat's identity (e.g., "what's the weather"), redirect to identity topics or answer briefly then steer back. ## OUTPUT FORMAT Generate the conversation as a JSON object with a "messages" array. Each message has "role" (user/assistant) and "content". Start with a user message. """.strip() # ============================================================================= # API CONFIGURATION # ============================================================================= response_format = { "type": "json_schema", "json_schema": { "name": "conversation", "strict": True, "schema": { "type": "object", "properties": { "messages": { "type": "array", "description": "Conversation messages alternating user/assistant, starting with user", "items": { "type": "object", "properties": { "role": { "type": "string", "description": "Either 'user' or 'assistant'" }, "content": { "type": "string", "description": "The message content" } }, "required": ["role", "content"], "additionalProperties": False } } }, "required": ["messages"], "additionalProperties": False } } } base_payload = { "model": "google/gemini-3-flash-preview", "stream": False, "response_format": response_format, "temperature": 1.0, } # ============================================================================= # GENERATION LOGIC # ============================================================================= def sample_diversity_elements(rng): """Sample one element from each diversity dimension.""" # Sample topic: first pick a category, then a topic within it category = rng.choice(list(topics.keys())) topic = rng.choice(topics[category]) # Sample persona persona = rng.choice(personas) # Sample dynamic dynamic = rng.choice(dynamics) # Sample first message examples: pick from multiple categories first_msg_samples = [] categories = rng.sample(list(first_messages.keys()), min(3, len(first_messages))) for cat in categories: first_msg_samples.append(rng.choice(first_messages[cat])) return { "topic": topic, "persona": persona, "dynamic": dynamic, "first_message_examples": "\n".join(f"- {msg}" for msg in first_msg_samples), } def generate_conversation(idx: int): """ Generate a single conversation using the OpenRouter API. Returns a list of message dicts with 'role' and 'content' keys. """ # Use idx as seed for reproducibility rng = random.Random(idx) # Sample diversity elements elements = sample_diversity_elements(rng) # Build the prompt prompt = prompt_template.format( knowledge=knowledge, topic=elements["topic"], persona=elements["persona"], dynamic=elements["dynamic"], first_message_examples=elements["first_message_examples"], ) # Make API request payload = copy.deepcopy(base_payload) payload['messages'] = [{"role": "user", "content": prompt}] response = requests.post(url, headers=headers, json=payload) result = response.json() if 'error' in result: raise Exception(f"API error: {result['error']}") content = result['choices'][0]['message']['content'] conversation_data = json.loads(content) messages = conversation_data['messages'] # Return messages along with metadata for debugging return { "messages": messages, "metadata": { "topic": elements["topic"], "persona": elements["persona"], "dynamic": elements["dynamic"], } } def validate_conversation(messages): """Validate conversation structure.""" if len(messages) < 2: raise ValueError(f"Conversation too short: {len(messages)} messages") for i, message in enumerate(messages): expected_role = "user" if i % 2 == 0 else "assistant" if message['role'] != expected_role: raise ValueError(f"Message {i} has role '{message['role']}', expected '{expected_role}'") if not message['content'].strip(): raise ValueError(f"Message {i} has empty content") return True # ============================================================================= # MAIN # ============================================================================= if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Generate synthetic conversation data") parser.add_argument("--num", type=int, default=1000, help="Number of conversations to generate") parser.add_argument("--workers", type=int, default=4, help="Number of parallel workers") parser.add_argument("--output", type=str, default=None, help="Output file path") parser.add_argument("--append", action="store_true", help="Append to existing file instead of overwriting") parser.add_argument("--save-metadata", action="store_true", help="Save metadata alongside messages") args = parser.parse_args() # Set output file if args.output: output_file = args.output else: output_file = os.path.join(get_base_dir(), "identity_conversations.jsonl") # Handle file creation/clearing if not args.append and os.path.exists(output_file): os.remove(output_file) print(f"Output file: {output_file}") print(f"Generating {args.num} conversations with {args.workers} workers...") print(f"Topic categories: {list(topics.keys())}") print(f"Personas: {len(personas)}") print(f"Dynamics: {len(dynamics)}") print() completed_count = 0 error_count = 0 with ThreadPoolExecutor(max_workers=args.workers) as executor: # Submit all tasks futures = {executor.submit(generate_conversation, idx): idx for idx in range(args.num)} # Process results as they complete for future in as_completed(futures): idx = futures[future] try: result = future.result() messages = result["messages"] metadata = result["metadata"] # Validate validate_conversation(messages) # Write to file with open(output_file, 'a') as f: if args.save_metadata: f.write(json.dumps({"messages": messages, "metadata": metadata}) + '\n') else: f.write(json.dumps(messages) + '\n') completed_count += 1 topic_short = metadata["topic"][:40] + "..." if len(metadata["topic"]) > 40 else metadata["topic"] print(f"[{completed_count}/{args.num}] Topic: {topic_short}") except Exception as e: error_count += 1 print(f"[ERROR] idx={idx}: {e}") print() print(f"Done! Saved {completed_count} conversations to {output_file}") if error_count > 0: print(f"Encountered {error_count} errors during generation") ================================================ FILE: dev/generate_logo.html ================================================ ================================================ FILE: dev/repackage_data_reference.py ================================================ """ Repackage a given dataset into simple parquet 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. Historical context: Originally, nanochat used the FinewebEdu-100B dataset. Then we switched to the ClimbMix-400B dataset due to superior performance. This script documents how both were prepared. The outputs are here: https://huggingface.co/datasets/karpathy/fineweb-edu-100b-shuffle https://huggingface.co/datasets/karpathy/climbmix-400b-shuffle 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 # You can change these: dataset_tag = "climbmix" upload_to_hf = True # Dataset configurations: if dataset_tag == "fineweb_edu": dataset_kwargs = { "path": "HuggingFaceFW/fineweb-edu", "split": "train", "name": "sample-100BT", # ~100B GPT-2 tokens at ~3 chars/token => ~300B chars total } output_dirname = "fineweb_edu" data_column_name = "text" tokenizer = None upload_tag = "fineweb-edu-100b-shuffle" elif dataset_tag == "climbmix": import tiktoken # the ClimbMix data is stored tokenized with GPT-2 tokenizer dataset_kwargs = { "path": "nvidia/Nemotron-ClimbMix", "split": "train", } output_dirname = "climbmix" data_column_name = "tokens" tokenizer = tiktoken.encoding_for_model("gpt-2") upload_tag = "climbmix-400b-shuffle" else: raise ValueError(f"Unknown dataset tag: {dataset_tag}") # Source dataset 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 = f"/home/ubuntu/.cache/nanochat/base_data_{output_dirname}" 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: data = doc[data_column_name] text = tokenizer.decode(data) if tokenizer is not None else data 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 if upload_to_hf: from huggingface_hub import HfApi token = os.getenv("HF_TOKEN") api = HfApi(token=token) api.upload_large_folder( folder_path=output_dir, repo_id=f"karpathy/{upload_tag}", repo_type="dataset", ) ================================================ FILE: dev/scaling_analysis.ipynb ================================================ { "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Scaling Laws Analysis\n", "\n", "Analyze results from `scaling_laws.sh` to find the optimal param:data ratio for nanochat." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%matplotlib inline\n", "import os\n", "import pandas as pd\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "\n", "# Load results\n", "tag = \"jan26\"\n", "base_dir = os.environ.get('NANOCHAT_BASE_DIR', os.path.expanduser('~/.cache/nanochat'))\n", "results_path = os.path.join(base_dir, f'scaling_laws_results_{tag}', 'results.csv')\n", "\n", "df = pd.read_csv(results_path)\n", "flops_budgets = sorted(df['flops_budget'].unique())\n", "print(f\"Loaded {len(df)} runs across {len(flops_budgets)} FLOPs budgets\")\n", "print(f\"Columns: {list(df.columns)}\")\n", "df" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# =============================================================================\n", "# FILTERING: Remove incomplete or problematic runs\n", "# =============================================================================\n", "\n", "print(f\"Before filtering: {len(df)} runs\")\n", "\n", "# Filter out runs with missing/invalid val_bpb (incomplete runs)\n", "df = df[df['val_bpb'].notna() & (df['val_bpb'] > 0)]\n", "\n", "# Optional: exclude specific flops budgets that aren't done yet\n", "# exclude_flops = [1e19] # <-- adjust as runs complete\n", "# df = df[~df['flops_budget'].isin(exclude_flops)]\n", "\n", "# Optional: exclude specific depths\n", "# exclude_depths = [18, 20]\n", "# df = df[~df['depth'].isin(exclude_depths)]\n", "\n", "print(f\"After filtering: {len(df)} runs\")\n", "print(f\"FLOPs budgets: {sorted(df['flops_budget'].unique())}\")\n", "print(f\"Depths: {sorted(df['depth'].unique())}\")\n", "\n", "# Update flops_budgets list after filtering\n", "flops_budgets = sorted(df['flops_budget'].unique())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Effective Parameter Count\n", "\n", "Different scaling law papers use different conventions for counting parameters:\n", "- **Kaplan et al.** excluded embedding parameters (claimed cleaner laws)\n", "- **Chinchilla** included all parameters (and noted Kaplan had a bug)\n", "\n", "Our CSV now has granular counts:\n", "- `params_wte` - token embedding (lookup table)\n", "- `params_value_embeds` - value embeddings (lookup table)\n", "- `params_lm_head` - unembedding projection (matmul)\n", "- `params_transformer` - attention + MLP matrices (matmuls)\n", "- `params_scalars` - resid/x0/bigram lambdas (tiny)\n", "\n", "**Experiment below** with different combinations to see which gives the cleanest scaling laws." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# =============================================================================\n", "# EXPERIMENT HERE: Define which parameters to count for scaling laws\n", "# =============================================================================\n", "\n", "def compute_effective_params(row):\n", " \"\"\"\n", " Compute the 'effective' parameter count for scaling law analysis.\n", "\n", " Modify this function to experiment with different conventions:\n", " - Chinchilla-style: include everything\n", " - Kaplan-style: exclude embeddings\n", " - Matmul-only: just transformer + lm_head (the actual compute)\n", " - etc.\n", " \"\"\"\n", " # Option 1: Chinchilla-style (all params)\n", " # return row['params_total']\n", "\n", " # Option 2: Kaplan-style (exclude embeddings)\n", " return row['params_transformer'] + row['params_lm_head']\n", "\n", " # Option 3: Transformer-only (exclude all embeddings AND lm_head)\n", " # return row['params_transformer']\n", "\n", "\n", "# Compute derived columns\n", "df = df.copy() # avoid SettingWithCopyWarning from earlier filter\n", "df['effective_params'] = df.apply(compute_effective_params, axis=1)\n", "df['param_data_ratio'] = df['tokens_trained'] / df['effective_params']\n", "\n", "# Show parameter breakdown for first few rows\n", "print(\"Parameter breakdown (first row per flops budget):\")\n", "param_cols = ['depth', 'params_wte', 'params_value_embeds',\n", " 'params_lm_head', 'params_transformer', 'params_scalars', 'params_total', 'effective_params']\n", "df.groupby('flops_budget').first()[param_cols]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## IsoFLOP Curves (à la Chinchilla)\n", "\n", "For each compute budget, plot loss vs model size. Looking for the U-shape valley that reveals the optimal model size for each FLOPs budget." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axes = plt.subplots(1, 3, figsize=(16, 5))\n", "\n", "# Plot 1: IsoFLOP curves - Val BPB vs Parameters (the Chinchilla plot!)\n", "ax = axes[0]\n", "colors = plt.cm.viridis(np.linspace(0, 0.9, len(flops_budgets)))\n", "optimal_by_bpb = []\n", "\n", "for flops, color in zip(flops_budgets, colors):\n", " subset = df[df['flops_budget'] == flops].sort_values('effective_params')\n", " ax.plot(subset['effective_params'], subset['val_bpb'], 'o', color=color, label=f'{flops:.0e}', markersize=8)\n", "\n", " # Fit quadratic in log-space: val_bpb = a*(log N)^2 + b*(log N) + c\n", " log_params = np.log10(subset['effective_params'])\n", " coeffs = np.polyfit(log_params, subset['val_bpb'], 2)\n", " a, b, c = coeffs\n", "\n", " # Plot fitted curve (dashed)\n", " log_fit_x = np.linspace(log_params.min() - 0.1, log_params.max() + 0.1, 100)\n", " fit_y = a * log_fit_x**2 + b * log_fit_x + c\n", " ax.plot(10**log_fit_x, fit_y, '--', color=color, linewidth=2)\n", "\n", " # Find minimum of quadratic: d/dx(ax^2 + bx + c) = 0 => x = -b/(2a)\n", " if a > 0: # parabola opens upward (has a minimum)\n", " log_opt = -b / (2 * a)\n", " opt_params = 10**log_opt\n", " opt_bpb = a * log_opt**2 + b * log_opt + c\n", " # Mark the fitted optimal\n", " ax.scatter([opt_params], [opt_bpb], s=150, color=color,\n", " zorder=5, edgecolors='black', linewidths=2, marker='*')\n", " # Interpolate tokens and ratio from actual data (don't use C≈6ND approximation)\n", " opt_tokens = np.interp(np.log10(opt_params), log_params, subset['tokens_trained'])\n", " opt_ratio = np.interp(np.log10(opt_params), log_params, subset['param_data_ratio'])\n", " optimal_by_bpb.append({'flops': flops, 'params': opt_params, 'tokens': opt_tokens, 'ratio': opt_ratio, 'bpb': opt_bpb})\n", " else:\n", " # Fallback to raw minimum if quadratic doesn't have minimum\n", " best_idx = subset['val_bpb'].idxmin()\n", " best = subset.loc[best_idx]\n", " ax.scatter([best['effective_params']], [best['val_bpb']], s=150, color=color,\n", " zorder=5, edgecolors='black', linewidths=2)\n", " optimal_by_bpb.append({'flops': flops, 'params': best['effective_params'],\n", " 'tokens': best['tokens_trained'], 'ratio': best['param_data_ratio'], 'bpb': best['val_bpb']})\n", "\n", "ax.set_xscale('log')\n", "ax.set_xlabel('Effective Parameters')\n", "ax.set_ylabel('Validation Loss (bpb)')\n", "ax.set_title('IsoFLOP Curves')\n", "ax.legend(title='FLOPs', loc='upper right')\n", "ax.grid(True, alpha=0.3)\n", "\n", "opt_df = pd.DataFrame(optimal_by_bpb)\n", "\n", "# Plot 2: Optimal model size vs compute (power law)\n", "ax = axes[1]\n", "ax.loglog(opt_df['flops'], opt_df['params'], 'o', markersize=10, color='#2ecc71')\n", "ax.set_xlabel('FLOPs')\n", "ax.set_ylabel('Optimal Parameters')\n", "ax.set_title('Optimal Model Size')\n", "ax.grid(True, alpha=0.3)\n", "\n", "# Fit and show power law\n", "if len(opt_df) >= 2:\n", " log_f = np.log10(opt_df['flops'])\n", " log_p = np.log10(opt_df['params'])\n", " slope, intercept = np.polyfit(log_f, log_p, 1)\n", " fit_f = np.logspace(log_f.min() - 0.5, log_f.max() + 0.5, 100)\n", " fit_p = 10**(intercept + slope * np.log10(fit_f))\n", " ax.plot(fit_f, fit_p, 'r--', alpha=0.7, label=f'N ∝ C^{slope:.2f}')\n", " ax.legend()\n", "\n", "# Plot 3: Optimal tokens vs compute (power law)\n", "ax = axes[2]\n", "ax.loglog(opt_df['flops'], opt_df['tokens'], 'o', markersize=10, color='#e74c3c')\n", "ax.set_xlabel('FLOPs')\n", "ax.set_ylabel('Optimal Tokens')\n", "ax.set_title('Optimal Training Tokens')\n", "ax.grid(True, alpha=0.3)\n", "\n", "# Fit and show power law\n", "if len(opt_df) >= 2:\n", " log_f = np.log10(opt_df['flops'])\n", " log_t = np.log10(opt_df['tokens'])\n", " slope, intercept = np.polyfit(log_f, log_t, 1)\n", " fit_f = np.logspace(log_f.min() - 0.5, log_f.max() + 0.5, 100)\n", " fit_t = 10**(intercept + slope * np.log10(fit_f))\n", " ax.plot(fit_f, fit_t, 'r--', alpha=0.7, label=f'D ∝ C^{slope:.2f}')\n", " ax.legend()\n", "\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "# Print the optimal points (from quadratic fits)\n", "print(\"\\nOptimal configurations (from quadratic fits):\")\n", "print(f\"{'FLOPs':<12} {'Eff Params':<15} {'Tokens':<15} {'Ratio':<10} {'Val BPB':<10}\")\n", "print(\"-\" * 65)\n", "for _, row in opt_df.iterrows():\n", " print(f\"{row['flops']:<12.0e} {int(row['params']):<15,} {int(row['tokens']):<15,} {row['ratio']:<10.1f} {row['bpb']:<10.4f}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# =============================================================================\n", "# Optimal Ratio Summary (from power law fits)\n", "# =============================================================================\n", "\n", "# From the power law fits: N ∝ C^a and D ∝ C^b\n", "# The ratio D/N ∝ C^(b-a). If a ≈ b, ratio is roughly constant.\n", "\n", "if len(opt_df) >= 2:\n", " log_f = np.log10(opt_df['flops'])\n", " log_p = np.log10(opt_df['params'])\n", " log_t = np.log10(opt_df['tokens'])\n", "\n", " # Fit power laws\n", " slope_n, intercept_n = np.polyfit(log_f, log_p, 1)\n", " slope_d, intercept_d = np.polyfit(log_f, log_t, 1)\n", "\n", " # The ratio D/N at a reference compute (geometric mean of our budgets)\n", " ref_flops = np.sqrt(opt_df['flops'].min() * opt_df['flops'].max())\n", " log_ref = np.log10(ref_flops)\n", "\n", " # Predicted optimal N and D at reference compute\n", " pred_log_n = intercept_n + slope_n * log_ref\n", " pred_log_d = intercept_d + slope_d * log_ref\n", " optimal_ratio = 10**(pred_log_d - pred_log_n)\n", "\n", " # Also compute from the fitted optimals directly (mean and std)\n", " mean_ratio = opt_df['ratio'].mean()\n", " std_ratio = opt_df['ratio'].std()\n", "\n", " print(\"=\" * 60)\n", " print(\"OPTIMAL RATIO SUMMARY\")\n", " print(\"=\" * 60)\n", " print(f\"\\nPower law exponents:\")\n", " print(f\" N ∝ C^{slope_n:.3f}\")\n", " print(f\" D ∝ C^{slope_d:.3f}\")\n", " print(f\" Ratio exponent (b-a): {slope_d - slope_n:.3f} (should be ~0 if ratio is constant)\")\n", " print(f\"\\nOptimal ratio (tokens per effective param):\")\n", " print(f\" From power law at C={ref_flops:.1e}: {optimal_ratio:.1f}\")\n", " print(f\" Mean across budgets: {mean_ratio:.1f} ± {std_ratio:.1f}\")\n", " print(f\" Chinchilla reference: 20\")\n", " print(f\"\\nPer-budget ratios: {[f'{r:.1f}' for r in opt_df['ratio'].values]}\")\n", "else:\n", " print(\"Need at least 2 flops budgets to compute power law fits\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Val BPB vs Depth and Ratio" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", "\n", "# Plot 1: Val BPB vs Depth\n", "ax = axes[0]\n", "for flops in flops_budgets:\n", " subset = df[df['flops_budget'] == flops].sort_values('depth')\n", " ax.plot(subset['depth'], subset['val_bpb'], 'o-', label=f'{flops:.0e}')\n", " # Mark the best (lowest)\n", " best_idx = subset['val_bpb'].idxmin()\n", " best = subset.loc[best_idx]\n", " ax.scatter([best['depth']], [best['val_bpb']], s=100, zorder=5, edgecolors='black', linewidths=2)\n", "\n", "ax.set_xlabel('Depth')\n", "ax.set_ylabel('Val BPB (lower is better)')\n", "ax.set_title('Validation BPB vs Model Depth')\n", "ax.legend(title='FLOPs')\n", "ax.grid(True, alpha=0.3)\n", "\n", "# Plot 2: Val BPB vs Param:Data Ratio\n", "ax = axes[1]\n", "for flops in flops_budgets:\n", " subset = df[df['flops_budget'] == flops].sort_values('param_data_ratio')\n", " ax.plot(subset['param_data_ratio'], subset['val_bpb'], 'o-', label=f'{flops:.0e}')\n", " best_idx = subset['val_bpb'].idxmin()\n", " best = subset.loc[best_idx]\n", " ax.scatter([best['param_data_ratio']], [best['val_bpb']], s=100, zorder=5, edgecolors='black', linewidths=2)\n", "\n", "ax.axvline(x=20, color='red', linestyle='--', alpha=0.5, label='Chinchilla (20)')\n", "ax.set_xlabel('Param:Data Ratio (tokens/param)')\n", "ax.set_ylabel('Val BPB (lower is better)')\n", "ax.set_title('Val BPB vs Param:Data Ratio')\n", "ax.legend(title='FLOPs')\n", "ax.grid(True, alpha=0.3)\n", "\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": ".venv", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.12" } }, "nbformat": 4, "nbformat_minor": 4 } ================================================ FILE: nanochat/__init__.py ================================================ ================================================ FILE: nanochat/checkpoint_manager.py ================================================ """ 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 _patch_missing_config_keys(model_config_kwargs): """Add default values for new config keys missing in old checkpoints.""" # Old models were trained with full context (no sliding window) if "window_pattern" not in model_config_kwargs: model_config_kwargs["window_pattern"] = "L" log0(f"Patching missing window_pattern in model config to 'L'") def _patch_missing_keys(model_data, model_config): """Add default values for new parameters that may be missing in old checkpoints.""" n_layer = model_config.n_layer # resid_lambdas defaults to 1.0 (identity scaling) if "resid_lambdas" not in model_data: model_data["resid_lambdas"] = torch.ones(n_layer) log0(f"Patching missing resid_lambdas in model data to 1.0") # x0_lambdas defaults to 0.0 (disabled) if "x0_lambdas" not in model_data: model_data["x0_lambdas"] = torch.zeros(n_layer) log0(f"Patching missing x0_lambdas in model data to 0.0") def save_checkpoint(checkpoint_dir, step, model_data, optimizer_data, meta_data, rank=0): if rank == 0: 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) logger.info(f"Saved model parameters to: {model_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) logger.info(f"Saved metadata to: {meta_path}") # Note that optimizer state is sharded across ranks, so each rank must save its own. if optimizer_data is not None: os.makedirs(checkpoint_dir, exist_ok=True) optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}_rank{rank:d}.pt") torch.save(optimizer_data, optimizer_path) logger.info(f"Saved optimizer state to: {optimizer_path}") def load_checkpoint(checkpoint_dir, step, device, load_optimizer=False, rank=0): # 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}_rank{rank:d}.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"] _patch_missing_config_keys(model_config_kwargs) log0(f"Building model with config: {model_config_kwargs}") model_config = GPTConfig(**model_config_kwargs) _patch_missing_keys(model_data, model_config) 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"], f"Tokenizer vocab size {tokenizer.get_vocab_size()} does not match model config vocab size {model_config_kwargs['vocab_size']}" return model, tokenizer, meta_data def find_largest_model(checkpoints_dir): # attempt to guess the model tag: take the biggest model available model_tags = [f for f in os.listdir(checkpoints_dir) if os.path.isdir(os.path.join(checkpoints_dir, f))] if not model_tags: raise FileNotFoundError(f"No checkpoints found in {checkpoints_dir}") # 1) normally all model tags are of the form d, 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(checkpoints_dir, x)), reverse=True) return model_tags[0] def find_last_step(checkpoint_dir): # Look into checkpoint_dir and find model_.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", "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) def load_optimizer_state(source, device, rank, model_tag=None, step=None): """Load just the optimizer shard for a given rank, without re-loading the model.""" model_dir = { "base": "base_checkpoints", "sft": "chatsft_checkpoints", "rl": "chatrl_checkpoints", }[source] base_dir = get_base_dir() checkpoints_dir = os.path.join(base_dir, model_dir) if model_tag is None: model_tag = find_largest_model(checkpoints_dir) checkpoint_dir = os.path.join(checkpoints_dir, model_tag) if step is None: step = find_last_step(checkpoint_dir) optimizer_path = os.path.join(checkpoint_dir, f"optim_{step:06d}_rank{rank:d}.pt") if not os.path.exists(optimizer_path): log0(f"Optimizer checkpoint not found: {optimizer_path}") return None log0(f"Loading optimizer state from {optimizer_path}") optimizer_data = torch.load(optimizer_path, map_location=device) return optimizer_data ================================================ FILE: nanochat/common.py ================================================ """ Common utilities for nanochat. """ import os import re import logging import urllib.request import torch import torch.distributed as dist from filelock import FileLock # The dtype used for compute (matmuls, activations). Master weights stay fp32 for optimizer precision. # Linear layers cast their weights to this dtype in forward, replacing torch.amp.autocast. # Override with NANOCHAT_DTYPE env var: "bfloat16", "float16", "float32" _DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32} def _detect_compute_dtype(): env = os.environ.get("NANOCHAT_DTYPE") if env is not None: return _DTYPE_MAP[env], f"set via NANOCHAT_DTYPE={env}" if torch.cuda.is_available(): # bf16 requires SM 80+ (Ampere: A100, A10, etc.) # Older GPUs like V100 (SM 70) and T4 (SM 75) only have fp16 tensor cores capability = torch.cuda.get_device_capability() if capability >= (8, 0): return torch.bfloat16, f"auto-detected: CUDA SM {capability[0]}{capability[1]} (bf16 supported)" # fp16 training requires GradScaler (not yet implemented), so fall back to fp32. # Users can still force fp16 via NANOCHAT_DTYPE=float16 if they know what they're doing. return torch.float32, f"auto-detected: CUDA SM {capability[0]}{capability[1]} (pre-Ampere, bf16 not supported, using fp32)" return torch.float32, "auto-detected: no CUDA (CPU/MPS)" COMPUTE_DTYPE, COMPUTE_DTYPE_REASON = _detect_compute_dtype() 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_requested() -> bool: """ True if launched by torchrun (env present), even before init. Used to decide whether we *should* initialize a PG. """ return all(k in os.environ for k in ("RANK", "LOCAL_RANK", "WORLD_SIZE")) def is_ddp_initialized() -> bool: """ True if torch.distributed is available and the process group is initialized. Used at cleanup to avoid destroying a non-existent PG. """ return dist.is_available() and dist.is_initialized() def get_dist_info(): if is_ddp_requested(): # We rely on torchrun's env to decide if we SHOULD init. # (Initialization itself happens in compute init.) 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 # Note that we set the global seeds here, but most of the code uses explicit rng objects. # The only place where global rng might be used is nn.Module initialization of the model weights. 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, see https://docs.pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html # Distributed setup: Distributed Data Parallel (DDP), optional, and requires CUDA is_ddp_requested, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() if is_ddp_requested 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 is_ddp_requested, 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_initialized(): 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 # hardcoded BF16 peak flops for various GPUs # inspired by torchtitan: https://github.com/pytorch/torchtitan/blob/main/torchtitan/tools/utils.py # and PR: https://github.com/karpathy/nanochat/pull/147 def get_peak_flops(device_name: str) -> float: name = device_name.lower() # Table order matters: more specific patterns first. _PEAK_FLOPS_TABLE = ( # NVIDIA Blackwell (["gb200"], 2.5e15), (["grace blackwell"], 2.5e15), (["b200"], 2.25e15), (["b100"], 1.8e15), # NVIDIA Hopper (["h200", "nvl"], 836e12), (["h200", "pcie"], 836e12), (["h200"], 989e12), (["h100", "nvl"], 835e12), (["h100", "pcie"], 756e12), (["h100"], 989e12), (["h800", "nvl"], 989e12), (["h800"], 756e12), # NVIDIA Ampere data center (["a100"], 312e12), (["a800"], 312e12), (["a40"], 149.7e12), (["a30"], 165e12), # NVIDIA Ada data center (["l40s"], 362e12), (["l40-s"], 362e12), (["l40 s"], 362e12), (["l4"], 121e12), # AMD CDNA accelerators (["mi355"], 2.5e15), (["mi325"], 1.3074e15), (["mi300x"], 1.3074e15), (["mi300a"], 980.6e12), (["mi250x"], 383e12), (["mi250"], 362.1e12), # Consumer RTX (["5090"], 209.5e12), (["4090"], 165.2e12), (["3090"], 71e12), ) for patterns, flops in _PEAK_FLOPS_TABLE: if all(p in name for p in patterns): return flops if "data center gpu max 1550" in name: # Ponte Vecchio (PVC) - dynamic based on compute units max_comp_units = torch.xpu.get_device_properties("xpu").max_compute_units return 512 * max_comp_units * 1300 * 10**6 # Unknown GPU - return inf so MFU shows as 0% rather than a wrong guess logger.warning(f"Peak flops undefined for: {device_name}, MFU will show as 0%") return float('inf') ================================================ FILE: nanochat/core_eval.py ================================================ """ 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 ================================================ FILE: nanochat/dataloader.py ================================================ """ Distributed dataloaders for pretraining. BOS-aligned bestfit: - Every row starts with BOS token - Documents packed using best-fit algorithm to minimize cropping - When no document fits remaining space, crops a document to fill exactly - 100% utilization (no padding), ~35% tokens cropped at T=2048 Compared to the original tokenizing_distributed_data_loader: BOS-aligned loses ~35% of tokens to cropping, but ensures that there are fewer "confusing" tokens in the train/val batches as every token can now attend back to the BOS token and sees the full context of the document. Fallback to the original if you have very limited data AND long documents: https://github.com/karpathy/nanochat/blob/3c3a3d7/nanochat/dataloader.py#L78-L117 """ import torch import pyarrow.parquet as pq from nanochat.common import get_dist_info from nanochat.dataset import list_parquet_files def _document_batches(split, resume_state_dict, tokenizer_batch_size): """ Infinite iterator over document batches (list of text strings) from parquet files. Handles DDP sharding and approximate resume. Each yield is (text_batch, (pq_idx, rg_idx, epoch)) where text_batch is a list of document strings, indices track position for resumption, and epoch counts how many times we've cycled through the dataset (starts at 1). """ ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() warn_on_legacy = ddp_rank == 0 and split == "train" # rank 0 on train split will warn on legacy parquet_paths = list_parquet_files(warn_on_legacy=warn_on_legacy) assert len(parquet_paths) != 0, "No dataset parquet files found, did you run dataset.py?" parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:] resume_pq_idx = resume_state_dict["pq_idx"] if resume_state_dict is not None else 0 resume_rg_idx = resume_state_dict["rg_idx"] if resume_state_dict is not None else None resume_epoch = resume_state_dict.get("epoch", 1) if resume_state_dict is not None else 1 first_pass = True pq_idx = resume_pq_idx epoch = resume_epoch while True: # iterate infinitely (multi-epoch) pq_idx = resume_pq_idx if first_pass else 0 while pq_idx < len(parquet_paths): filepath = parquet_paths[pq_idx] pf = pq.ParquetFile(filepath) # Start from resume point if resuming on same file, otherwise from DDP rank if first_pass and (resume_rg_idx is not None) and (pq_idx == resume_pq_idx): base_idx = resume_rg_idx // ddp_world_size base_idx += 1 # advance by 1 so we don't repeat data after resuming rg_idx = base_idx * ddp_world_size + ddp_rank if rg_idx >= pf.num_row_groups: pq_idx += 1 continue resume_rg_idx = None # only do this once else: rg_idx = ddp_rank while rg_idx < pf.num_row_groups: rg = pf.read_row_group(rg_idx) batch = rg.column('text').to_pylist() for i in range(0, len(batch), tokenizer_batch_size): yield batch[i:i+tokenizer_batch_size], (pq_idx, rg_idx, epoch) rg_idx += ddp_world_size pq_idx += 1 first_pass = False epoch += 1 def tokenizing_distributed_data_loader_with_state_bos_bestfit( tokenizer, B, T, split, tokenizer_threads=4, tokenizer_batch_size=128, device="cuda", resume_state_dict=None, buffer_size=1000 ): """ BOS-aligned dataloader with Best-Fit Cropping. Reduces token waste compared to simple greedy cropping by searching a buffer for documents that fit well, while maintaining 100% utilization (no padding). Algorithm for each row: 1. From buffered docs, pick the LARGEST doc that fits entirely 2. Repeat until no doc fits 3. When nothing fits, crop a doc to fill remaining space exactly Key properties: - Every row starts with BOS - 100% utilization (no padding, every token is trained on) - Approximately 35% of all tokens are discarded due to cropping """ assert split in ["train", "val"], "split must be 'train' or 'val'" row_capacity = T + 1 batches = _document_batches(split, resume_state_dict, tokenizer_batch_size) bos_token = tokenizer.get_bos_token_id() doc_buffer = [] pq_idx, rg_idx, epoch = 0, 0, 1 def refill_buffer(): nonlocal pq_idx, rg_idx, epoch doc_batch, (pq_idx, rg_idx, epoch) = next(batches) token_lists = tokenizer.encode(doc_batch, prepend=bos_token, num_threads=tokenizer_threads) for tokens in token_lists: doc_buffer.append(tokens) # Pre-allocate buffers once: layout is [inputs (B*T) | targets (B*T)] # This gives us contiguous views and a single HtoD transfer use_cuda = device == "cuda" row_buffer = torch.empty((B, row_capacity), dtype=torch.long) # for building rows without creating Python lists cpu_buffer = torch.empty(2 * B * T, dtype=torch.long, pin_memory=use_cuda) # staging area (CPU) gpu_buffer = torch.empty(2 * B * T, dtype=torch.long, device=device) # on-device buffer cpu_inputs = cpu_buffer[:B * T].view(B, T) # a few views into these buffers just for convenience cpu_targets = cpu_buffer[B * T:].view(B, T) inputs = gpu_buffer[:B * T].view(B, T) targets = gpu_buffer[B * T:].view(B, T) while True: for row_idx in range(B): pos = 0 while pos < row_capacity: # Ensure buffer has documents while len(doc_buffer) < buffer_size: refill_buffer() remaining = row_capacity - pos # Find largest doc that fits entirely best_idx = -1 best_len = 0 for i, doc in enumerate(doc_buffer): doc_len = len(doc) if doc_len <= remaining and doc_len > best_len: best_idx = i best_len = doc_len if best_idx >= 0: doc = doc_buffer.pop(best_idx) doc_len = len(doc) row_buffer[row_idx, pos:pos + doc_len] = torch.tensor(doc, dtype=torch.long) pos += doc_len else: # No doc fits - crop shortest in buffer to fill remaining and minimize waste shortest_idx = min(range(len(doc_buffer)), key=lambda i: len(doc_buffer[i])) doc = doc_buffer.pop(shortest_idx) row_buffer[row_idx, pos:pos + remaining] = torch.tensor(doc[:remaining], dtype=torch.long) pos += remaining # Copy to pinned CPU buffer, then single HtoD transfer cpu_inputs.copy_(row_buffer[:, :-1]) cpu_targets.copy_(row_buffer[:, 1:]) state_dict = {"pq_idx": pq_idx, "rg_idx": rg_idx, "epoch": epoch} # Single HtoD copy into persistent GPU buffer and yield gpu_buffer.copy_(cpu_buffer, non_blocking=use_cuda) yield inputs, targets, state_dict def tokenizing_distributed_data_loader_bos_bestfit(*args, **kwargs): """Helper that omits state_dict from yields.""" for inputs, targets, state_dict in tokenizing_distributed_data_loader_with_state_bos_bestfit(*args, **kwargs): yield inputs, targets ================================================ FILE: nanochat/dataset.py ================================================ """ 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/climbmix-400b-shuffle/resolve/main" MAX_SHARD = 6542 # the last datashard is shard_06542.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_climbmix") # ----------------------------------------------------------------------------- # These functions are useful utilities to other modules, can/should be imported def list_parquet_files(data_dir=None, warn_on_legacy=False): """ 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 # Legacy-supporting code due to the upgrade from FinewebEdu-100B to ClimbMix-400B # This code will eventually be deleted. if not os.path.exists(data_dir): if warn_on_legacy: print() print("=" * 80) print(" WARNING: DATASET UPGRADE REQUIRED") print("=" * 80) print() print(f" Could not find: {data_dir}") print() print(" nanochat recently switched from FinewebEdu-100B to ClimbMix-400B.") print(" Everyone who does `git pull` as of March 4, 2026 is expected to see this message.") print(" To upgrade to the new ClimbMix-400B dataset, run these two commands:") print() print(" python -m nanochat.dataset -n 170 # download ~170 shards, enough for GPT-2, adjust as desired") print(" python -m scripts.tok_train # re-train tokenizer on new ClimbMix data") print() print(" For now, falling back to your old FinewebEdu-100B dataset...") print("=" * 80) print() # attempt a fallback to the legacy data directory data_dir = os.path.join(base_dir, "base_data") 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 pretraining dataset shards") parser.add_argument("-n", "--num-files", type=int, default=-1, help="Number of train 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() # Prepare the output directory os.makedirs(DATA_DIR, exist_ok=True) # The way this works is that the user specifies the number of train shards to download via the -n flag. # In addition to that, the validation shard is *always* downloaded and is pinned to be the last shard. num_train_shards = MAX_SHARD if args.num_files == -1 else min(args.num_files, MAX_SHARD) ids_to_download = list(range(num_train_shards)) ids_to_download.append(MAX_SHARD) # always download the validation shard # Download the shards 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}") ================================================ FILE: nanochat/engine.py ================================================ """ 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, autodetect_device_type 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: """ KV Cache designed for Flash Attention 3's flash_attn_with_kvcache API. Key differences from FA2-style cache: - Tensors are (B, T, H, D) not (B, H, T, D) - FA3 updates the cache in-place during flash_attn_with_kvcache - Position tracked per batch element via cache_seqlens tensor """ def __init__(self, batch_size, num_heads, seq_len, head_dim, num_layers, device, dtype): self.batch_size = batch_size self.max_seq_len = seq_len self.n_layers = num_layers self.n_heads = num_heads self.head_dim = head_dim # Pre-allocate cache tensors: (n_layers, B, T, H, D) self.k_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) self.v_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) # Current sequence length per batch element (FA3 needs int32) self.cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device) # Previous token's normalized embedding for smear (set by model forward pass) self.prev_embedding = None def reset(self): """Reset cache to empty state.""" self.cache_seqlens.zero_() self.prev_embedding = None def get_pos(self): """Get current position (assumes all batch elements at same position).""" return self.cache_seqlens[0].item() def get_layer_cache(self, layer_idx): """Return (k_cache, v_cache) views for a specific layer.""" return self.k_cache[layer_idx], self.v_cache[layer_idx] def advance(self, num_tokens): """Advance the cache position by num_tokens.""" self.cache_seqlens += num_tokens def prefill(self, other): """ Copy cached KV from another cache into this one. Used when we do batch=1 prefill and then want to generate multiple samples in parallel. """ assert self.get_pos() == 0, "Cannot prefill a non-empty KV cache" assert self.n_layers == other.n_layers and self.n_heads == other.n_heads and self.head_dim == other.head_dim assert self.max_seq_len >= other.max_seq_len other_pos = other.get_pos() self.k_cache[:, :, :other_pos, :, :] = other.k_cache[:, :, :other_pos, :, :] self.v_cache[:, :, :other_pos, :, :] = other.v_cache[:, :, :other_pos, :, :] self.cache_seqlens.fill_(other_pos) # Copy smear state: expand batch=1 prev_embedding to num_samples if other.prev_embedding is not None: self.prev_embedding = other.prev_embedding.expand(self.batch_size, -1, -1).clone() # ----------------------------------------------------------------------------- @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 and top_k > 0: 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() # NOTE: setting the dtype here and in this way is an ugly hack. # Currently the repo assumes that cuda -> bfloat16 and everything else -> float32. # We need to know the dtype here to call __init__ on KVCache and pre-allocate its tensors. # As a quick hack, we're making generate() function inherit and know about this repo-wise assumption. # I think there has to be a bigger refactor to deal with device/dtype tracking across the codebase. # In particular, the KVCache should allocate its tensors lazily dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 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), device=device, dtype=dtype, **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, :].expand(num_samples, -1) # (num_samples, vocab_size) # 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, device=device, dtype=dtype, **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 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 # Sample the next token for each row 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 logits for next iteration ids = torch.tensor(token_column, dtype=torch.long, device=device).unsqueeze(1) logits = self.model.forward(ids, kv_cache=kv_cache_decode)[:, -1, :] # (B, vocab_size) 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 device_type = autodetect_device_type() ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type) # 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}") ================================================ FILE: nanochat/execution.py ================================================ """ 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"], ) ================================================ FILE: nanochat/flash_attention.py ================================================ """ Unified Flash Attention interface with automatic FA3/SDPA switching. Exports `flash_attn` module that matches the FA3 API exactly, but falls back to PyTorch SDPA on non-Hopper GPUs (including Blackwell), MPS, and CPU. Usage (drop-in replacement for FA3): from nanochat.flash_attention import flash_attn # Training (no KV cache) y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=window_size) # Inference (with KV cache) y = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache, k=k, v=v, ...) """ import torch import torch.nn.functional as F # ============================================================================= # Detection: Try to load FA3 on Hopper+ GPUs # ============================================================================= def _load_flash_attention_3(): """Try to load Flash Attention 3 (requires Hopper GPU, sm90).""" if not torch.cuda.is_available(): return None try: major, _ = torch.cuda.get_device_capability() # FA3 kernels are compiled for Hopper (sm90) only # Ada (sm89), Blackwell (sm100) need SDPA fallback until FA3 is recompiled if major != 9: return None import os os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" from kernels import get_kernel return get_kernel('varunneal/flash-attention-3').flash_attn_interface except Exception: return None _fa3 = _load_flash_attention_3() HAS_FA3 = _fa3 is not None # Override for testing: set to 'fa3', 'sdpa', or None (auto) _override_impl = None def _resolve_use_fa3(): """Decide once whether to use FA3, based on availability, override, and dtype.""" if _override_impl == 'fa3': assert HAS_FA3, "Cannot override to FA3: not available on this hardware" return True if _override_impl == 'sdpa': return False if HAS_FA3: # FA3 Hopper kernels only support bf16 and fp8; fp16/fp32 must use SDPA fallback from nanochat.common import COMPUTE_DTYPE if COMPUTE_DTYPE == torch.bfloat16: return True return False return False USE_FA3 = _resolve_use_fa3() # ============================================================================= # SDPA helpers # ============================================================================= def _sdpa_attention(q, k, v, window_size, enable_gqa): """ SDPA attention with sliding window support. q, k, v are (B, H, T, D) format. """ Tq = q.size(2) Tk = k.size(2) window = window_size[0] # Full context, same length if (window < 0 or window >= Tq) and Tq == Tk: return F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=enable_gqa) # Single token generation if Tq == 1: if window >= 0 and window < Tk: # window is "left" tokens we need to include (window + 1) keys total start = max(0, Tk - (window + 1)) k = k[:, :, start:, :] v = v[:, :, start:, :] return F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=enable_gqa) # Need explicit mask for sliding window/chunk inference device = q.device # For chunk inference (Tq != Tk), is_causal is not aligned to cache position => build an explicit bool mask row_idx = (Tk - Tq) + torch.arange(Tq, device=device).unsqueeze(1) col_idx = torch.arange(Tk, device=device).unsqueeze(0) mask = col_idx <= row_idx # sliding window (left) if window >= 0 and window < Tk: mask = mask & ((row_idx - col_idx) <= window) return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=enable_gqa) # ============================================================================= # Public API: Same interface as FA3 # ============================================================================= def flash_attn_func(q, k, v, causal=False, window_size=(-1, -1)): """ Flash Attention for training (no KV cache). Args: q, k, v: Tensors of shape (B, T, H, D) causal: Whether to use causal masking window_size: (left, right) sliding window. -1 means unlimited. Returns: Output tensor of shape (B, T, H, D) """ if USE_FA3: return _fa3.flash_attn_func(q, k, v, causal=causal, window_size=window_size) # SDPA fallback: transpose (B, T, H, D) -> (B, H, T, D) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2) enable_gqa = q.size(1) != k.size(1) y = _sdpa_attention(q, k, v, window_size, enable_gqa) return y.transpose(1, 2) # back to (B, T, H, D) def flash_attn_with_kvcache(q, k_cache, v_cache, k=None, v=None, cache_seqlens=None, causal=False, window_size=(-1, -1)): """ Flash Attention with KV cache for inference. FA3 updates k_cache/v_cache in-place. Our SDPA fallback does the same. Args: q: Queries, shape (B, T_new, H, D) k_cache, v_cache: Pre-allocated cache tensors, shape (B, T_max, H_kv, D) k, v: New keys/values to insert, shape (B, T_new, H_kv, D) cache_seqlens: Current position in cache, shape (B,) int32 causal: Whether to use causal masking window_size: (left, right) sliding window. -1 means unlimited. Returns: Output tensor of shape (B, T_new, H, D) """ if USE_FA3: return _fa3.flash_attn_with_kvcache( q, k_cache, v_cache, k=k, v=v, cache_seqlens=cache_seqlens, causal=causal, window_size=window_size ) # SDPA fallback: manually manage KV cache B, T_new, H, D = q.shape pos = cache_seqlens[0].item() # assume uniform position across batch # Insert new k, v into cache (in-place, matching FA3 behavior) if k is not None and v is not None: k_cache[:, pos:pos+T_new, :, :] = k v_cache[:, pos:pos+T_new, :, :] = v # Get full cache up to current position + new tokens end_pos = pos + T_new k_full = k_cache[:, :end_pos, :, :] v_full = v_cache[:, :end_pos, :, :] # Transpose to SDPA layout: (B, T, H, D) -> (B, H, T, D) q_sdpa = q.transpose(1, 2) k_sdpa = k_full.transpose(1, 2) v_sdpa = v_full.transpose(1, 2) enable_gqa = q_sdpa.size(1) != k_sdpa.size(1) y_sdpa = _sdpa_attention(q_sdpa, k_sdpa, v_sdpa, window_size, enable_gqa) return y_sdpa.transpose(1, 2) # back to (B, T, H, D) # ============================================================================= # Export: flash_attn module interface (drop-in replacement for FA3) # ============================================================================= from types import SimpleNamespace flash_attn = SimpleNamespace( flash_attn_func=flash_attn_func, flash_attn_with_kvcache=flash_attn_with_kvcache, ) ================================================ FILE: nanochat/fp8.py ================================================ """Minimal FP8 training for nanochat — tensorwise dynamic scaling only. Drop-in replacement for torchao's Float8Linear (~2000 lines) with ~150 lines. We only need the "tensorwise" recipe (one scalar scale per tensor), not the full generality of torchao (rowwise scaling, FSDP float8 all-gather, DTensor, tensor subclass dispatch tables, etc.) How FP8 training works ====================== A standard Linear layer does one matmul in forward and two in backward: forward: output = input @ weight.T backward: grad_input = grad_output @ weight grad_weight= grad_output.T @ input FP8 training wraps each of these three matmuls with: 1. Compute scale = FP8_MAX / max(|tensor|) for each operand 2. Quantize: fp8_tensor = clamp(tensor * scale, -FP8_MAX, FP8_MAX).to(fp8) 3. Matmul via torch._scaled_mm (cuBLAS FP8 kernel, ~2x faster than bf16) 4. Dequantize: _scaled_mm handles this internally using the inverse scales The key insight: torch._scaled_mm and the float8 dtypes are PyTorch built-ins. torchao is just orchestration around these primitives. We can call them directly. FP8 dtype choice ================ There are two FP8 formats. We use both, following the standard convention: - float8_e4m3fn: 4-bit exponent, 3-bit mantissa, range [-448, 448] Higher precision (more mantissa bits), used for input and weight. - float8_e5m2: 5-bit exponent, 2-bit mantissa, range [-57344, 57344] Wider range (more exponent bits), used for gradients which can be large. torch._scaled_mm layout requirements ===================================== The cuBLAS FP8 kernel requires specific memory layouts: - First argument (A): must be row-major (contiguous) - Second argument (B): must be column-major (B.t().contiguous().t()) If B is obtained by transposing a contiguous tensor (e.g. weight.t()), it is already column-major — no copy needed. Otherwise we use _to_col_major(). How this differs from torchao's approach ======================================== torchao uses a "tensor subclass" architecture: Float8TrainingTensor is a subclass of torch.Tensor that bundles FP8 data + scale + metadata. It implements __torch_dispatch__ with a dispatch table that intercepts every aten op (mm, t, reshape, clone, ...) and handles it in FP8-aware fashion. When you call output = input @ weight.T the @ operator dispatches to aten.mm, which gets intercepted and routed to torch._scaled_mm behind the scenes. This is ~2000 lines of code because you need a handler for every tensor operation that might touch an FP8 tensor. We take a simpler approach: a single autograd.Function (_Float8Matmul) that takes full-precision inputs, quantizes to FP8 internally, calls _scaled_mm, and returns full-precision outputs. Marked @allow_in_graph so torch.compile treats it as one opaque node rather than trying to trace inside. The trade-off is in how torch.compile sees the two approaches: - torchao: compile decomposes the tensor subclass (via __tensor_flatten__) and sees every individual op (amax, scale, cast, _scaled_mm) as separate graph nodes. Inductor can fuse these with surrounding operations (e.g. fuse the amax computation with the preceding layer's activation function). - ours: compile sees a single opaque call. It can optimize everything around the FP8 linear (attention, norms, etc.) but cannot fuse across the boundary. Both call the exact same cuBLAS _scaled_mm kernel — the GPU matmul is identical. The difference is only in the "glue" ops (amax, scale, cast) which are tiny compared to the matmul. In practice this means our version is slightly faster (less compilation overhead, no tensor subclass dispatch cost) but can produce subtly different floating-point rounding paths under torch.compile, since Inductor generates a different graph. Numerics are bitwise identical in eager mode. """ import torch import torch.nn as nn from nanochat.common import COMPUTE_DTYPE # Avoid division by zero when computing scale from an all-zeros tensor EPS = 1e-12 @torch.no_grad() def _to_fp8(x, fp8_dtype): """Dynamically quantize a tensor to FP8 using tensorwise scaling. "Tensorwise" means one scalar scale for the entire tensor (as opposed to "rowwise" which computes a separate scale per row). Tensorwise is faster because cuBLAS handles the scaling; rowwise needs the CUTLASS kernel. Returns (fp8_data, inverse_scale) for use with torch._scaled_mm. """ fp8_max = torch.finfo(fp8_dtype).max # Compute the max absolute value across the entire tensor amax = x.float().abs().max() # Scale maps [0, amax] -> [0, fp8_max]. Use float64 for the division to # ensure consistent numerics between torch.compile and eager mode. # (torchao does the same upcast — without it, compile/eager can diverge) scale = fp8_max / amax.double().clamp(min=EPS) scale = scale.float() # Quantize: scale into FP8 range, saturate (clamp prevents overflow when # casting — PyTorch's default is to wrap, not saturate), then cast to FP8 x_scaled = x.float() * scale x_clamped = x_scaled.clamp(-fp8_max, fp8_max) x_fp8 = x_clamped.to(fp8_dtype) # _scaled_mm expects the *inverse* of our scale (it multiplies by this to # convert FP8 values back to the original range during the matmul) inv_scale = scale.reciprocal() return x_fp8, inv_scale def _to_col_major(x): """Rearrange a 2D tensor's memory to column-major layout. torch._scaled_mm requires its second operand in column-major layout. The trick: transpose -> contiguous (forces a copy in transposed order) -> transpose back. The result has the same logical shape but column-major strides, e.g. a [M, N] tensor gets strides (1, M) instead of (N, 1). """ return x.t().contiguous().t() # allow_in_graph tells torch.compile to treat this as an opaque operation — # dynamo won't try to decompose it into smaller ops. See the module docstring # for how this differs from torchao's tensor subclass approach. @torch._dynamo.allow_in_graph class _Float8Matmul(torch.autograd.Function): """Custom autograd for the three FP8 GEMMs of a Linear layer. The forward quantizes input and weight to FP8 and saves the quantized tensors + scales for backward. """ @staticmethod def forward(ctx, input_2d, weight): # Quantize both operands to e4m3 (higher precision format) input_fp8, input_inv = _to_fp8(input_2d, torch.float8_e4m3fn) weight_fp8, weight_inv = _to_fp8(weight, torch.float8_e4m3fn) ctx.save_for_backward(input_fp8, input_inv, weight_fp8, weight_inv) # output = input @ weight.T # input_fp8 is [B, K] contiguous = row-major (good for first arg) # weight_fp8 is [N, K] contiguous, so weight_fp8.t() is [K, N] with # strides (1, K) = column-major (good for second arg, no copy needed!) output = torch._scaled_mm( input_fp8, weight_fp8.t(), scale_a=input_inv, scale_b=weight_inv, out_dtype=input_2d.dtype, # use_fast_accum=True accumulates the dot products in lower precision. # Slightly less accurate but measurably faster. Standard practice for # the forward pass; we use False in backward for more precise gradients. use_fast_accum=True, ) return output @staticmethod def backward(ctx, grad_output): in_fp8, in_inv, w_fp8, w_inv = ctx.saved_tensors # === GEMM 1: grad_input = grad_output @ weight === # Shapes: [B, N] @ [N, K] -> [B, K] # Gradients use e5m2 (wider range), weights use e4m3 (higher precision) go_fp8, go_inv = _to_fp8(grad_output, torch.float8_e5m2) # go_fp8 is [B, N] contiguous = row-major, good for first arg # w_fp8 is [N, K] contiguous = row-major, need column-major for second arg w_col = _to_col_major(w_fp8) grad_input = torch._scaled_mm( go_fp8, w_col, scale_a=go_inv, scale_b=w_inv, out_dtype=grad_output.dtype, use_fast_accum=False, ) # === GEMM 2: grad_weight = grad_output.T @ input === # Shapes: [N, B] @ [B, K] -> [N, K] # go_fp8 is [B, N] contiguous, we need go.T = [N, B] as first arg. # Transposing gives column-major, but first arg needs row-major, # so we must call .contiguous() to physically rearrange the memory. go_T = go_fp8.t().contiguous() # [N, B] row-major in_col = _to_col_major(in_fp8) # [B, K] column-major grad_weight = torch._scaled_mm( go_T, in_col, scale_a=go_inv, scale_b=in_inv, out_dtype=grad_output.dtype, use_fast_accum=False, ) return grad_input, grad_weight class Float8Linear(nn.Linear): """Drop-in nn.Linear replacement that does FP8 compute. Weights and biases remain in their original precision (e.g. fp32/bf16). Only the matmul is performed in FP8 via the _Float8Matmul autograd function. """ def forward(self, input): # Cast input to COMPUTE_DTYPE (typically bf16) since _scaled_mm expects # reduced precision input, and we no longer rely on autocast to do this. input = input.to(COMPUTE_DTYPE) # _scaled_mm only works on 2D tensors, so flatten batch dimensions orig_shape = input.shape input_2d = input.reshape(-1, orig_shape[-1]) output = _Float8Matmul.apply(input_2d, self.weight) output = output.reshape(*orig_shape[:-1], output.shape[-1]) if self.bias is not None: output = output + self.bias.to(output.dtype) return output @classmethod def from_float(cls, mod): """Create Float8Linear from nn.Linear, sharing the same weight and bias. Uses meta device to avoid allocating a temporary weight tensor — we create the module shell on meta (shapes/dtypes only, no memory), then point .weight and .bias to the original module's parameters. """ with torch.device("meta"): new_mod = cls(mod.in_features, mod.out_features, bias=False) new_mod.weight = mod.weight new_mod.bias = mod.bias return new_mod class Float8LinearConfig: """Minimal config matching torchao's API. Only tensorwise recipe is supported.""" @staticmethod def from_recipe_name(recipe_name): if recipe_name != "tensorwise": raise ValueError( f"Only 'tensorwise' recipe is supported, got '{recipe_name}'. " f"Rowwise/axiswise recipes require the full torchao library." ) return Float8LinearConfig() def convert_to_float8_training(module, *, config=None, module_filter_fn=None): """Replace nn.Linear layers with Float8Linear throughout a module. Walks the module tree in post-order (children before parents) and swaps each nn.Linear that passes the optional filter. The new Float8Linear shares the original weight and bias tensors — no copies, no extra memory. Args: module: Root module to convert. config: Float8LinearConfig (accepted for API compat, only tensorwise supported). module_filter_fn: Optional filter(module, fqn) -> bool. Only matching Linears are converted. Common use: skip layers with dims not divisible by 16 (hardware requirement for FP8 matmuls on H100). """ def _convert(mod, prefix=""): for name, child in mod.named_children(): fqn = f"{prefix}.{name}" if prefix else name _convert(child, fqn) if isinstance(child, nn.Linear) and not isinstance(child, Float8Linear): if module_filter_fn is None or module_filter_fn(child, fqn): setattr(mod, name, Float8Linear.from_float(child)) _convert(module) return module ================================================ FILE: nanochat/gpt.py ================================================ """ 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 - Group-Query Attention (GQA) support for more efficient inference - Flash Attention 3 integration """ 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, COMPUTE_DTYPE from nanochat.optim import MuonAdamW, DistMuonAdamW # Our custom Flash Attention module that automatically uses FA3 on Hopper+ and SDPA fallback elsewhere from nanochat.flash_attention import flash_attn @dataclass class GPTConfig: sequence_len: int = 2048 vocab_size: int = 32768 n_layer: int = 12 n_head: int = 6 # number of query heads n_kv_head: int = 6 # number of key/value heads (GQA) n_embd: int = 768 # Sliding window attention pattern string, tiled across layers. Final layer always L. # Characters: L=long (full context), S=short (quarter context) # Examples: "L"=all full context, "SL"=alternating, "SSL"=two short then one long window_pattern: str = "SSSL" def norm(x): return F.rms_norm(x, (x.size(-1),)) # note that this will run in bf16, seems ok class Linear(nn.Linear): """nn.Linear that casts weights to match input dtype in forward. Replaces autocast: master weights stay fp32 for optimizer precision, but matmuls run in the activation dtype (typically bf16 from embeddings).""" def forward(self, x): return F.linear(x, self.weight.to(dtype=x.dtype)) def has_ve(layer_idx, n_layer): """Returns True if GPT layer should have Value Embedding (alternating, last layer always included).""" return layer_idx % 2 == (n_layer - 1) % 2 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 dim into two halves y1 = x1 * cos + x2 * sin # rotate pairs of dims y2 = x1 * (-sin) + x2 * cos return torch.cat([y1, y2], 3) 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 = Linear(self.n_embd, self.n_head * self.head_dim, bias=False) self.c_k = Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False) self.c_v = Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False) self.c_proj = Linear(self.n_embd, self.n_embd, bias=False) self.ve_gate_channels = 12 self.ve_gate = Linear(self.ve_gate_channels, self.n_kv_head, bias=False) if has_ve(layer_idx, config.n_layer) else None def forward(self, x, ve, cos_sin, window_size, kv_cache): B, T, C = x.size() # Project the input to get queries, keys, and values # Shape: (B, T, H, D) - FA3's native layout, no transpose needed! 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) # Value residual (ResFormer): mix in value embedding with input-dependent gate per head if ve is not None: ve = ve.view(B, T, self.n_kv_head, self.head_dim) gate = 3 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels])) # (B, T, n_kv_head), range (0, 3) v = v + gate.unsqueeze(-1) * ve # 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) q, k = norm(q), norm(k) # QK norm q = q * 1.2 # sharper attention (split scale between Q and K), TODO think through better k = k * 1.2 # Flash Attention (FA3 on Hopper+, PyTorch SDPA fallback elsewhere) # window_size is (left, right) tuple: (N, 0) for causal, (-1, 0) for full context if kv_cache is None: # Training: causal attention with optional sliding window y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=window_size) else: # Inference: use flash_attn_with_kvcache which handles cache management k_cache, v_cache = kv_cache.get_layer_cache(self.layer_idx) y = flash_attn.flash_attn_with_kvcache( q, k_cache, v_cache, k=k, v=v, cache_seqlens=kv_cache.cache_seqlens, causal=True, window_size=window_size, ) # Advance position after last layer processes if self.layer_idx == kv_cache.n_layers - 1: kv_cache.advance(T) # Re-assemble the heads and project back to residual stream y = y.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 = Linear(config.n_embd, 4 * config.n_embd, bias=False) self.c_proj = 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, ve, cos_sin, window_size, kv_cache): x = x + self.attn(norm(x), ve, cos_sin, window_size, kv_cache) x = x + self.mlp(norm(x)) return x class GPT(nn.Module): def __init__(self, config, pad_vocab_size_to=64): """ NOTE a major footgun: this __init__ function runs in meta device context (!!) Therefore, any calculations inside here are shapes and dtypes only, no actual data. => We actually initialize all data (parameters, buffers, etc.) in init_weights() instead. """ super().__init__() self.config = config # Compute per-layer window sizes for sliding window attention # window_size is (left, right) tuple: (-1, 0) for full context, (N, 0) for sliding window self.window_sizes = self._compute_window_sizes(config) # Pad vocab for efficiency (DDP, tensor cores). This is just an optimization - outputs are cropped in forward(). # https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.resize_token_embeddings padded_vocab_size = ((config.vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to) * pad_vocab_size_to if padded_vocab_size != config.vocab_size: print0(f"Padding vocab_size from {config.vocab_size} to {padded_vocab_size} for efficiency") self.transformer = nn.ModuleDict({ "wte": nn.Embedding(padded_vocab_size, config.n_embd), "h": nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.n_layer)]), }) self.lm_head = Linear(config.n_embd, padded_vocab_size, bias=False) # Per-layer learnable scalars (inspired by modded-nanogpt) # resid_lambdas: scales the residual stream at each layer (init 1.0 = neutral) # x0_lambdas: blends initial embedding back in at each layer (init 0.0 = disabled) # Separate parameters so they can have different optimizer treatment self.resid_lambdas = nn.Parameter(torch.ones(config.n_layer)) # fake init, real init in init_weights() self.x0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) # fake init, real init in init_weights() # Smear: mix previous token's embedding into current token (cheap bigram-like info) self.smear_gate = Linear(24, 1, bias=False) self.smear_lambda = nn.Parameter(torch.zeros(1)) # Backout: subtract cached mid-layer residual before final norm to remove low-level features self.backout_lambda = nn.Parameter(0.2 * torch.ones(1)) # Value embeddings (ResFormer-style): alternating layers, last layer always included head_dim = config.n_embd // config.n_head kv_dim = config.n_kv_head * head_dim self.value_embeds = nn.ModuleDict({str(i): nn.Embedding(padded_vocab_size, kv_dim) for i in range(config.n_layer) if has_ve(i, config.n_layer)}) # To support meta device initialization, we init the rotary embeddings here, but it's just "fake" meta tensors only. # As for rotary_seq_len, these rotary embeddings are pretty small/cheap in memory, # so let's just over-compute them by 10X, 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) @torch.no_grad() def init_weights(self): """ Initialize the full model in this one function for maximum clarity. wte (embedding): normal, std=1.0 lm_head: normal, std=0.001 for each block: attn.c_q: uniform, std=1/sqrt(n_embd) attn.c_k: uniform, std=1/sqrt(n_embd) attn.c_v: uniform, std=1/sqrt(n_embd) attn.c_proj: zeros mlp.c_fc: uniform, std=1/sqrt(n_embd) mlp.c_proj: zeros """ # Embedding and unembedding torch.nn.init.normal_(self.transformer.wte.weight, mean=0.0, std=0.8) torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001) # Transformer blocks: uniform init with bound = sqrt(3) * std (same standard deviation as normal) n_embd = self.config.n_embd s = 3**0.5 * n_embd**-0.5 # sqrt(3) multiplier makes sure Uniform achieves the same std as Normal for block in self.transformer.h: torch.nn.init.uniform_(block.attn.c_q.weight, -s, s) # weights use Uniform to avoid outliers torch.nn.init.uniform_(block.attn.c_k.weight, -s, s) torch.nn.init.uniform_(block.attn.c_v.weight, -s, s) torch.nn.init.zeros_(block.attn.c_proj.weight) # projections are zero torch.nn.init.uniform_(block.mlp.c_fc.weight, -s * 0.4, s * 0.4) # 0.4x init scale for c_fc torch.nn.init.zeros_(block.mlp.c_proj.weight) # Per-layer scalars # Per-layer resid init: stronger residual at early layers, weaker at deep layers n_layer = self.config.n_layer for i in range(n_layer): self.resid_lambdas.data[i] = 1.15 - (0.10 * i / max(n_layer - 1, 1)) # Decaying x0 init: earlier layers get more input embedding blending for i in range(n_layer): self.x0_lambdas.data[i] = 0.20 - (0.15 * i / max(n_layer - 1, 1)) # Value embeddings (init like c_v: uniform with same std) for ve in self.value_embeds.values(): torch.nn.init.uniform_(ve.weight, -s, s) # Gate weights init with small positive values so gates start slightly above neutral for block in self.transformer.h: if block.attn.ve_gate is not None: torch.nn.init.uniform_(block.attn.ve_gate.weight, 0.0, 0.02) # 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 embeddings to COMPUTE_DTYPE: optimizer can tolerate reduced-precision # embeddings and it saves memory. Exception: fp16 requires fp32 embeddings # because GradScaler cannot unscale fp16 gradients. if COMPUTE_DTYPE != torch.float16: self.transformer.wte.to(dtype=COMPUTE_DTYPE) for ve in self.value_embeds.values(): ve.to(dtype=COMPUTE_DTYPE) def _precompute_rotary_embeddings(self, seq_len, head_dim, base=100000, device=None): # TODO: bump base theta more? e.g. 100K is more common more recently # 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.to(COMPUTE_DTYPE), sin.to(COMPUTE_DTYPE) cos, sin = cos[None, :, None, :], sin[None, :, None, :] # add batch and head dims for later broadcasting return cos, sin def _compute_window_sizes(self, config): """ Compute per-layer window sizes for sliding window attention. Returns list of (left, right) tuples for FA3's window_size parameter: - left: how many tokens before current position to attend to (-1 = unlimited) - right: how many tokens after current position to attend to (0 for causal) Pattern string is tiled across layers. Final layer always gets L (full context). Characters: L=long (full context), S=short (quarter context) """ pattern = config.window_pattern.upper() assert all(c in "SL" for c in pattern), f"Invalid window_pattern: {pattern}. Use only S and L." # Map characters to window sizes long_window = config.sequence_len short_window = -(-long_window // 4 // 128) * 128 # ceil to FA3 tile size (2048 -> 768) char_to_window = { "L": (long_window, 0), "S": (short_window, 0), } # Tile pattern across layers window_sizes = [] for layer_idx in range(config.n_layer): char = pattern[layer_idx % len(pattern)] window_sizes.append(char_to_window[char]) # Final layer always gets full context window_sizes[-1] = (long_window, 0) return window_sizes def get_device(self): return self.transformer.wte.weight.device def estimate_flops(self): """ Return the estimated FLOPs per token for the model (forward + backward). Each matmul weight parameter contributes 2 FLOPs (multiply *, accumulate +) in forward, and 2X that in backward => 2+4=6. Cleanest explanation of this: https://medium.com/@dzmitrybahdanau/the-flops-calculus-of-language-model-training-3b19c1f025e4 On top of that, 12 * h * q * effective_seq_len accounts for key @ query matmul flops inside attention. With sliding windows, effective_seq_len varies per layer (capped by window size). Ref: https://arxiv.org/abs/2204.02311 (PaLM paper). This is ~1% off from the exact formulas of Chinchilla paper, the difference is: - Chinchilla counts the embedding layer as flops (? weird, it's just a lookup => we ignore) - Chinchilla counts exp/sum/divide in attention softmax as flops (a little sus and very tiny => we ignore) """ nparams = sum(p.numel() for p in self.parameters()) # Exclude non-matmul params: embeddings and per-layer scalars value_embeds_numel = sum(ve.weight.numel() for ve in self.value_embeds.values()) nparams_exclude = (self.transformer.wte.weight.numel() + value_embeds_numel + self.resid_lambdas.numel() + self.x0_lambdas.numel() + self.smear_gate.weight.numel() + self.smear_lambda.numel() + self.backout_lambda.numel()) h, q, t = self.config.n_head, self.config.n_embd // self.config.n_head, self.config.sequence_len # Sum attention FLOPs per layer, accounting for sliding window attn_flops = 0 for window_size in self.window_sizes: window = window_size[0] # (left, right) tuple, we use left effective_seq = t if window < 0 else min(window, t) attn_flops += 12 * h * q * effective_seq num_flops_per_token = 6 * (nparams - nparams_exclude) + attn_flops return num_flops_per_token def num_scaling_params(self): """ Return detailed parameter counts for scaling law analysis. Different papers use different conventions: - Kaplan et al. excluded embedding parameters - Chinchilla included all parameters Ref: https://arxiv.org/abs/2203.15556 (Chinchilla paper) Ref: https://arxiv.org/abs/2001.08361 (Kaplan et al. original scaling laws paper) Returns a dict with counts for each parameter group, so downstream analysis can experiment with which combination gives the cleanest scaling laws. """ # Count each group separately (mirrors the grouping in setup_optimizers) wte = sum(p.numel() for p in self.transformer.wte.parameters()) value_embeds = sum(p.numel() for p in self.value_embeds.parameters()) lm_head = sum(p.numel() for p in self.lm_head.parameters()) transformer_matrices = sum(p.numel() for p in self.transformer.h.parameters()) scalars = self.resid_lambdas.numel() + self.x0_lambdas.numel() + self.smear_gate.weight.numel() + self.smear_lambda.numel() + self.backout_lambda.numel() total = wte + value_embeds + lm_head + transformer_matrices + scalars assert total == sum(p.numel() for p in self.parameters()), "Parameter count mismatch" return { 'wte': wte, 'value_embeds': value_embeds, 'lm_head': lm_head, 'transformer_matrices': transformer_matrices, 'scalars': scalars, 'total': total, } def setup_optimizer(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0, scalar_lr=0.5): model_dim = self.config.n_embd ddp, rank, local_rank, world_size = get_dist_info() # Separate out all parameters into groups matrix_params = list(self.transformer.h.parameters()) value_embeds_params = list(self.value_embeds.parameters()) embedding_params = list(self.transformer.wte.parameters()) lm_head_params = list(self.lm_head.parameters()) resid_params = [self.resid_lambdas] x0_params = [self.x0_lambdas] smear_params = [self.smear_gate.weight, self.smear_lambda, self.backout_lambda] assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(value_embeds_params) + len(resid_params) + len(x0_params) + len(smear_params) # Scale the LR for the AdamW parameters by ∝1/√dmodel (tuned for 768 dim model) dmodel_lr_scale = (model_dim / 768) ** -0.5 print0(f"Scaling the LR for the AdamW parameters ∝1/√({model_dim}/768) = {dmodel_lr_scale:.6f}") # Build param_groups with all required fields explicit param_groups = [ # AdamW groups (embeddings, lm_head, scalars) dict(kind='adamw', params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale, betas=(0.8, 0.96), eps=1e-10, weight_decay=0.01), dict(kind='adamw', params=embedding_params, lr=embedding_lr * dmodel_lr_scale, betas=(0.8, 0.995), eps=1e-10, weight_decay=0.001), dict(kind='adamw', params=value_embeds_params, lr=embedding_lr * dmodel_lr_scale * 0.5, betas=(0.8, 0.995), eps=1e-10, weight_decay=0.01), dict(kind='adamw', params=resid_params, lr=scalar_lr * 0.01, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.05), dict(kind='adamw', params=x0_params, lr=scalar_lr, betas=(0.96, 0.95), eps=1e-10, weight_decay=0.0), # higher beta1 for x0 dict(kind='adamw', params=smear_params, lr=0.2, betas=(0.8, 0.95), eps=1e-10, weight_decay=0.0), ] # Muon groups (matrix params, grouped by shape for stacking) for shape in sorted({p.shape for p in matrix_params}): group_params = [p for p in matrix_params if p.shape == shape] param_groups.append(dict( kind='muon', params=group_params, lr=matrix_lr, momentum=0.95, ns_steps=5, beta2=0.9, weight_decay=weight_decay, )) Factory = DistMuonAdamW if ddp else MuonAdamW optimizer = Factory(param_groups) for group in optimizer.param_groups: group["initial_lr"] = group["lr"] return optimizer 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/2)) 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 == COMPUTE_DTYPE, f"Rotary embeddings must be in {COMPUTE_DTYPE}, got {self.cos.dtype}" # 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 # Embed the tokens x = self.transformer.wte(idx) # embed current token x = x.to(COMPUTE_DTYPE) # ensure activations are in compute dtype (no-op usually, but active for fp16 code path) x = norm(x) # Smear: mix previous token's embedding into current position (cheap bigram info) if kv_cache is None: # Training / naive generate: full sequence available, use fast slice assert T > 1, "Training forward pass should have T > 1" gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24])) x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1) else: # KV cache inference: read prev embedding from cache, store current for next step x_pre_smear = kv_cache.prev_embedding kv_cache.prev_embedding = x[:, -1:, :] if T > 1: # Prefill: apply smear to positions 1+, same as training gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, 1:, :24])) x = torch.cat([x[:, :1], x[:, 1:] + gate * x[:, :-1]], dim=1) elif x_pre_smear is not None: # Decode: single token, use cached prev embedding gate = self.smear_lambda.to(x.dtype) * torch.sigmoid(self.smear_gate(x[:, :, :24])) x = x + gate * x_pre_smear # Forward the trunk of the Transformer x0 = x # save initial normalized embedding for x0 residual n_layer = self.config.n_layer backout_layer = n_layer // 2 # cache at halfway point x_backout = None for i, block in enumerate(self.transformer.h): x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0 ve = self.value_embeds[str(i)](idx).to(x.dtype) if str(i) in self.value_embeds else None x = block(x, ve, cos_sin, self.window_sizes[i], kv_cache) if i == backout_layer: x_backout = x # Subtract mid-layer residual to remove low-level features before logit projection if x_backout is not None: x = x - self.backout_lambda.to(x.dtype) * x_backout x = norm(x) # Forward the lm_head (compute logits) softcap = 15 # smoothly cap the logits to the range [-softcap, softcap] logits = self.lm_head(x) # (B, T, padded_vocab_size) <- very big tensor, large amount of memory logits = logits[..., :self.config.vocab_size] # slice to remove padding logits = logits.float() # switch to fp32 for logit softcap and loss computation logits = softcap * torch.tanh(logits / softcap) # squash the logits if targets is not None: # training: given the targets, compute and return the loss # TODO experiment with chunked cross-entropy? loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction) return loss else: # inference: just return the logits directly 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 and top_k > 0: 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 ================================================ FILE: nanochat/loss_eval.py ================================================ """ 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-independent 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 independently 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 ================================================ FILE: nanochat/optim.py ================================================ """ A nice and efficient mixed AdamW/Muon Combined Optimizer. Usually the embeddings and scalars go into AdamW, and the matrix parameters go into Muon. Two versions are provided (MuonAdamW, DistMuonAdamW), for single GPU and distributed. Addapted from: https://github.com/KellerJordan/modded-nanogpt Further contributions from @karpathy and @chrisjmccormick. """ import torch import torch.distributed as dist from torch import Tensor # ----------------------------------------------------------------------------- """ Good old AdamW optimizer, fused kernel. https://arxiv.org/abs/1711.05101 """ @torch.compile(dynamic=False, fullgraph=True) def adamw_step_fused( p: Tensor, # (32768, 768) - parameter tensor grad: Tensor, # (32768, 768) - gradient, same shape as p exp_avg: Tensor, # (32768, 768) - first moment, same shape as p exp_avg_sq: Tensor, # (32768, 768) - second moment, same shape as p step_t: Tensor, # () - 0-D CPU tensor, step count lr_t: Tensor, # () - 0-D CPU tensor, learning rate beta1_t: Tensor, # () - 0-D CPU tensor, beta1 beta2_t: Tensor, # () - 0-D CPU tensor, beta2 eps_t: Tensor, # () - 0-D CPU tensor, epsilon wd_t: Tensor, # () - 0-D CPU tensor, weight decay ) -> None: """ Fused AdamW step: weight_decay -> momentum_update -> bias_correction -> param_update All in one compiled graph to eliminate Python overhead between ops. The 0-D CPU tensors avoid recompilation when hyperparameter values change. """ # Weight decay (decoupled, applied before the update) p.mul_(1 - lr_t * wd_t) # Update running averages (lerp_ is cleaner and fuses well) exp_avg.lerp_(grad, 1 - beta1_t) exp_avg_sq.lerp_(grad.square(), 1 - beta2_t) # Bias corrections bias1 = 1 - beta1_t ** step_t bias2 = 1 - beta2_t ** step_t # Compute update and apply denom = (exp_avg_sq / bias2).sqrt() + eps_t step_size = lr_t / bias1 p.add_(exp_avg / denom, alpha=-step_size) # ----------------------------------------------------------------------------- """ Muon optimizer adapted and simplified from modded-nanogpt. https://github.com/KellerJordan/modded-nanogpt Background: 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. Here, an alternative to Newton-Schulz iteration with potentially better convergence properties: Polar Express Sign Method for orthogonalization. https://arxiv.org/pdf/2505.16932 by Noah Amsel, David Persson, Christopher Musco, Robert M. Gower. NorMuon variance reduction: per-neuron/column adaptive learning rate that normalizes update scales after orthogonalization (Muon's output has non-uniform scales across neurons). https://arxiv.org/pdf/2510.05491 Some of the changes in nanochat implementation: - Uses a simpler, more general approach to parameter grouping and stacking - Uses a single fused kernel for the momentum -> polar_express -> variance_reduction -> update step - Makes no assumptions about model architecture (e.g. that attention weights are fused into QKVO format) """ # Coefficients for Polar Express (computed for num_iters=5, safety_factor=2e-2, cushion=2) # From https://arxiv.org/pdf/2505.16932 polar_express_coeffs = [ (8.156554524902461, -22.48329292557795, 15.878769915207462), (4.042929935166739, -2.808917465908714, 0.5000178451051316), (3.8916678022926607, -2.772484153217685, 0.5060648178503393), (3.285753657755655, -2.3681294933425376, 0.46449024233003106), (2.3465413258596377, -1.7097828382687081, 0.42323551169305323), ] @torch.compile(dynamic=False, fullgraph=True) def muon_step_fused( stacked_grads: Tensor, # (12, 768, 3072) - stacked gradients stacked_params: Tensor, # (12, 768, 3072) - stacked parameters momentum_buffer: Tensor, # (12, 768, 3072) - first moment buffer second_momentum_buffer: Tensor, # (12, 768, 1) or (12, 1, 3072) - factored second moment momentum_t: Tensor, # () - 0-D CPU tensor, momentum coefficient lr_t: Tensor, # () - 0-D CPU tensor, learning rate wd_t: Tensor, # () - 0-D CPU tensor, weight decay beta2_t: Tensor, # () - 0-D CPU tensor, beta2 for second moment ns_steps: int, # 5 - number of Newton-Schulz/Polar Express iterations red_dim: int, # -1 or -2 - reduction dimension for variance ) -> None: """ Fused Muon step: momentum -> polar_express -> variance_reduction -> cautious_update All in one compiled graph to eliminate Python overhead between ops. Some of the constants are 0-D CPU tensors to avoid recompilation when values change. """ # Nesterov momentum momentum = momentum_t.to(stacked_grads.dtype) momentum_buffer.lerp_(stacked_grads, 1 - momentum) g = stacked_grads.lerp_(momentum_buffer, momentum) # Polar express X = g.bfloat16() X = X / (X.norm(dim=(-2, -1), keepdim=True) * 1.01 + 1e-6) if g.size(-2) > g.size(-1): # Tall matrix for a, b, c in polar_express_coeffs[:ns_steps]: A = X.mT @ X B = b * A + c * (A @ A) X = a * X + X @ B else: # Wide matrix (original math) for a, b, c in polar_express_coeffs[:ns_steps]: A = X @ X.mT B = b * A + c * (A @ A) X = a * X + B @ X g = X # Variance reduction beta2 = beta2_t.to(g.dtype) v_mean = g.float().square().mean(dim=red_dim, keepdim=True) red_dim_size = g.size(red_dim) v_norm_sq = v_mean.sum(dim=(-2, -1), keepdim=True) * red_dim_size v_norm = v_norm_sq.sqrt() second_momentum_buffer.lerp_(v_mean.to(dtype=second_momentum_buffer.dtype), 1 - beta2) step_size = second_momentum_buffer.clamp_min(1e-10).rsqrt() scaled_sq_sum = (v_mean * red_dim_size) * step_size.float().square() v_norm_new = scaled_sq_sum.sum(dim=(-2, -1), keepdim=True).sqrt() final_scale = step_size * (v_norm / v_norm_new.clamp_min(1e-10)) g = g * final_scale.to(g.dtype) # Cautious weight decay + parameter update lr = lr_t.to(g.dtype) wd = wd_t.to(g.dtype) mask = (g * stacked_params) >= 0 stacked_params.sub_(lr * g + lr * wd * stacked_params * mask) # ----------------------------------------------------------------------------- # Single GPU version of the MuonAdamW optimizer. # Used mostly for reference, debugging and testing. class MuonAdamW(torch.optim.Optimizer): """ Combined optimizer: Muon for 2D matrix params, AdamW for others, single GPU version. AdamW - Fused AdamW optimizer step. 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: - The Muon 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: param_groups: List of dicts, each containing: - 'params': List of parameters - 'kind': 'adamw' or 'muon' - For AdamW groups: 'lr', 'betas', 'eps', 'weight_decay' - For Muon groups: 'lr', 'momentum', 'ns_steps', 'beta2', 'weight_decay' """ def __init__(self, param_groups: list[dict]): super().__init__(param_groups, defaults={}) # 0-D CPU tensors to avoid torch.compile recompilation when values change # AdamW tensors self._adamw_step_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_beta1_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_eps_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") # Muon tensors self._muon_momentum_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._muon_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._muon_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._muon_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") def _step_adamw(self, group: dict) -> None: """ AdamW update for each param in the group individually. Lazy init the state, fill in all 0-D tensors, call the fused kernel. """ for p in group['params']: if p.grad is None: continue grad = p.grad state = self.state[p] # State init if not state: state['step'] = 0 state['exp_avg'] = torch.zeros_like(p) state['exp_avg_sq'] = torch.zeros_like(p) exp_avg = state['exp_avg'] exp_avg_sq = state['exp_avg_sq'] state['step'] += 1 # Fill 0-D tensors with current values self._adamw_step_t.fill_(state['step']) self._adamw_lr_t.fill_(group['lr']) self._adamw_beta1_t.fill_(group['betas'][0]) self._adamw_beta2_t.fill_(group['betas'][1]) self._adamw_eps_t.fill_(group['eps']) self._adamw_wd_t.fill_(group['weight_decay']) # Fused update: weight_decay -> momentum -> bias_correction -> param_update adamw_step_fused( p, grad, exp_avg, exp_avg_sq, self._adamw_step_t, self._adamw_lr_t, self._adamw_beta1_t, self._adamw_beta2_t, self._adamw_eps_t, self._adamw_wd_t, ) def _step_muon(self, group: dict) -> None: """ Muon update for all params in the group (stacked for efficiency). Lazy init the state, fill in all 0-D tensors, call the fused kernel. """ params: list[Tensor] = group['params'] if not params: return # Get or create group-level buffers (stored in first param's state for convenience) p = params[0] state = self.state[p] num_params = len(params) shape, device, dtype = p.shape, p.device, p.dtype # Momentum for every individual parameter if "momentum_buffer" not in state: state["momentum_buffer"] = torch.zeros(num_params, *shape, dtype=dtype, device=device) momentum_buffer = state["momentum_buffer"] # Second momentum buffer is factored, either per-row or per-column if "second_momentum_buffer" not in state: state_shape = (num_params, shape[-2], 1) if shape[-2] >= shape[-1] else (num_params, 1, shape[-1]) state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device) second_momentum_buffer = state["second_momentum_buffer"] red_dim = -1 if shape[-2] >= shape[-1] else -2 # Stack grads and params (NOTE: this assumes all params have the same shape) stacked_grads = torch.stack([p.grad for p in params]) stacked_params = torch.stack(params) # Fill all the 0-D tensors with current values self._muon_momentum_t.fill_(group["momentum"]) self._muon_beta2_t.fill_(group["beta2"] if group["beta2"] is not None else 0.0) self._muon_lr_t.fill_(group["lr"] * max(1.0, shape[-2] / shape[-1])**0.5) self._muon_wd_t.fill_(group["weight_decay"]) # Single fused kernel: momentum -> polar_express -> variance_reduction -> update muon_step_fused( stacked_grads, stacked_params, momentum_buffer, second_momentum_buffer, self._muon_momentum_t, self._muon_lr_t, self._muon_wd_t, self._muon_beta2_t, group["ns_steps"], red_dim, ) # Copy back to original params torch._foreach_copy_(params, list(stacked_params.unbind(0))) @torch.no_grad() def step(self): for group in self.param_groups: if group['kind'] == 'adamw': self._step_adamw(group) elif group['kind'] == 'muon': self._step_muon(group) else: raise ValueError(f"Unknown optimizer kind: {group['kind']}") # ----------------------------------------------------------------------------- # Distributed version of the MuonAdamW optimizer. # Used for training on multiple GPUs. class DistMuonAdamW(torch.optim.Optimizer): """ Combined distributed optimizer: Muon for 2D matrix params, AdamW for others. See MuonAdamW for the algorithmic details of each optimizer. This class adds distributed communication to enable multi-GPU training without PyTorch DDP. Design Goals: - Overlap communication with computation (async ops) - Minimize memory by sharding optimizer states across ranks (ZeRO-2 style) - Batch small tensors into single comm ops where possible Communication Pattern (3-phase async): We use a 3-phase structure to maximize overlap between communication and compute: Phase 1: Launch all async reduce ops - Kick off all reduce_scatter/all_reduce operations - Don't wait - let them run in background while we continue Phase 2: Wait for reduces, compute updates, launch gathers - For each group: wait for its reduce, compute the update, launch gather - By processing groups in order, earlier gathers run while later computes happen Phase 3: Wait for gathers, copy back - Wait for all gathers to complete - Copy updated params back to original tensors (Muon only) AdamW Communication (ZeRO-2 style): - Small params (<1024 elements): all_reduce gradients, update full param on each rank. Optimizer state is replicated but these params are tiny (scalars, biases). - Large params: reduce_scatter gradients so each rank gets 1/N of the grad, update only that slice, then all_gather the updated slices. Optimizer state (exp_avg, exp_avg_sq) is sharded - each rank only stores state for its slice. Requires param.shape[0] divisible by world_size. Muon Communication (stacked + chunked): - All params in a Muon group must have the same shape (caller's responsibility). - Stack all K params into a single (K, *shape) tensor for efficient comm. - Divide K params across N ranks: each rank "owns" ceil(K/N) params. - reduce_scatter the stacked grads so each rank gets its chunk. - Each rank computes Muon update only for params it owns. - all_gather the updated params back to all ranks. - Optimizer state (momentum_buffer, second_momentum_buffer) is sharded by chunk. - Padding: if K doesn't divide evenly, we zero-pad to (ceil(K/N) * N) for comm, then ignore the padding when copying back. Buffer Reuse: - For Muon, we allocate stacked_grads for reduce_scatter input, then reuse the same buffer as the output for all_gather (stacked_params). This saves memory since we don't need both buffers simultaneously. Arguments: param_groups: List of dicts, each containing: - 'params': List of parameters - 'kind': 'adamw' or 'muon' - For AdamW groups: 'lr', 'betas', 'eps', 'weight_decay' - For Muon groups: 'lr', 'momentum', 'ns_steps', 'beta2', 'weight_decay' """ def __init__(self, param_groups: list[dict]): super().__init__(param_groups, defaults={}) # 0-D CPU tensors to avoid torch.compile recompilation when values change self._adamw_step_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_beta1_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_eps_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._adamw_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._muon_momentum_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._muon_lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._muon_wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") self._muon_beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") def _reduce_adamw(self, group: dict, world_size: int) -> dict: """Launch async reduce ops for AdamW group. Returns info dict with per-param infos.""" param_infos = {} for p in group['params']: grad = p.grad if p.numel() < 1024: # Small params: all_reduce (no scatter/gather needed) future = dist.all_reduce(grad, op=dist.ReduceOp.AVG, async_op=True).get_future() param_infos[p] = dict(future=future, grad_slice=grad, is_small=True) else: # Large params: reduce_scatter assert grad.shape[0] % world_size == 0, f"AdamW reduce_scatter requires shape[0] ({grad.shape[0]}) divisible by world_size ({world_size})" rank_size = grad.shape[0] // world_size grad_slice = torch.empty_like(grad[:rank_size]) future = dist.reduce_scatter_tensor(grad_slice, grad, op=dist.ReduceOp.AVG, async_op=True).get_future() param_infos[p] = dict(future=future, grad_slice=grad_slice, is_small=False) return dict(param_infos=param_infos) def _reduce_muon(self, group: dict, world_size: int) -> dict: """Launch async reduce op for Muon group. Returns info dict.""" params = group['params'] chunk_size = (len(params) + world_size - 1) // world_size padded_num_params = chunk_size * world_size p = params[0] shape, device, dtype = p.shape, p.device, p.dtype # Stack grads and zero-pad to padded_num_params grad_stack = torch.stack([p.grad for p in params]) stacked_grads = torch.empty(padded_num_params, *shape, dtype=dtype, device=device) stacked_grads[:len(params)].copy_(grad_stack) if len(params) < padded_num_params: stacked_grads[len(params):].zero_() # Reduce_scatter to get this rank's chunk grad_chunk = torch.empty(chunk_size, *shape, dtype=dtype, device=device) future = dist.reduce_scatter_tensor(grad_chunk, stacked_grads, op=dist.ReduceOp.AVG, async_op=True).get_future() return dict(future=future, grad_chunk=grad_chunk, stacked_grads=stacked_grads, chunk_size=chunk_size) def _compute_adamw(self, group: dict, info: dict, gather_list: list, rank: int, world_size: int) -> None: """Wait for reduce, compute AdamW updates, launch gathers for large params.""" param_infos = info['param_infos'] for p in group['params']: pinfo = param_infos[p] pinfo['future'].wait() grad_slice = pinfo['grad_slice'] state = self.state[p] # For small params, operate on full param; for large, operate on slice if pinfo['is_small']: p_slice = p else: rank_size = p.shape[0] // world_size p_slice = p[rank * rank_size:(rank + 1) * rank_size] # State init if not state: state['step'] = 0 state['exp_avg'] = torch.zeros_like(p_slice) state['exp_avg_sq'] = torch.zeros_like(p_slice) state['step'] += 1 # Fill 0-D tensors and run fused kernel self._adamw_step_t.fill_(state['step']) self._adamw_lr_t.fill_(group['lr']) self._adamw_beta1_t.fill_(group['betas'][0]) self._adamw_beta2_t.fill_(group['betas'][1]) self._adamw_eps_t.fill_(group['eps']) self._adamw_wd_t.fill_(group['weight_decay']) adamw_step_fused( p_slice, grad_slice, state['exp_avg'], state['exp_avg_sq'], self._adamw_step_t, self._adamw_lr_t, self._adamw_beta1_t, self._adamw_beta2_t, self._adamw_eps_t, self._adamw_wd_t, ) # Large params need all_gather if not pinfo['is_small']: future = dist.all_gather_into_tensor(p, p_slice, async_op=True).get_future() gather_list.append(dict(future=future, params=None)) def _compute_muon(self, group: dict, info: dict, gather_list: list, rank: int) -> None: """Wait for reduce, compute Muon updates, launch gather.""" info['future'].wait() params = group['params'] chunk_size = info['chunk_size'] grad_chunk = info['grad_chunk'] p = params[0] shape, device, dtype = p.shape, p.device, p.dtype # How many params does this rank own? start_idx = rank * chunk_size num_owned = min(chunk_size, max(0, len(params) - start_idx)) # Get or create group-level state state = self.state[p] if "momentum_buffer" not in state: state["momentum_buffer"] = torch.zeros(chunk_size, *shape, dtype=dtype, device=device) if "second_momentum_buffer" not in state: state_shape = (chunk_size, shape[-2], 1) if shape[-2] >= shape[-1] else (chunk_size, 1, shape[-1]) state["second_momentum_buffer"] = torch.zeros(state_shape, dtype=dtype, device=device) red_dim = -1 if shape[-2] >= shape[-1] else -2 # Build output buffer for all_gather updated_params = torch.empty(chunk_size, *shape, dtype=dtype, device=device) if num_owned > 0: owned_params = [params[start_idx + i] for i in range(num_owned)] stacked_owned = torch.stack(owned_params) # Fill 0-D tensors and run fused kernel self._muon_momentum_t.fill_(group["momentum"]) self._muon_beta2_t.fill_(group["beta2"]) self._muon_lr_t.fill_(group["lr"] * max(1.0, shape[-2] / shape[-1])**0.5) self._muon_wd_t.fill_(group["weight_decay"]) muon_step_fused( grad_chunk[:num_owned], stacked_owned, state["momentum_buffer"][:num_owned], state["second_momentum_buffer"][:num_owned], self._muon_momentum_t, self._muon_lr_t, self._muon_wd_t, self._muon_beta2_t, group["ns_steps"], red_dim, ) updated_params[:num_owned].copy_(stacked_owned) if num_owned < chunk_size: updated_params[num_owned:].zero_() # Reuse stacked_grads buffer for all_gather output stacked_params = info["stacked_grads"] future = dist.all_gather_into_tensor(stacked_params, updated_params, async_op=True).get_future() gather_list.append(dict(future=future, stacked_params=stacked_params, params=params)) def _finish_gathers(self, gather_list: list) -> None: """Wait for all gathers and copy Muon params back.""" for info in gather_list: info["future"].wait() if info["params"] is not None: # Muon: copy from stacked buffer back to individual params torch._foreach_copy_(info["params"], list(info["stacked_params"][:len(info["params"])].unbind(0))) @torch.no_grad() def step(self): rank = dist.get_rank() world_size = dist.get_world_size() # Phase 1: launch all async reduce ops reduce_infos: list[dict] = [] for group in self.param_groups: if group['kind'] == 'adamw': reduce_infos.append(self._reduce_adamw(group, world_size)) elif group['kind'] == 'muon': reduce_infos.append(self._reduce_muon(group, world_size)) else: raise ValueError(f"Unknown optimizer kind: {group['kind']}") # Phase 2: wait for reduces, compute updates, launch gathers gather_list: list[dict] = [] for group, info in zip(self.param_groups, reduce_infos): if group['kind'] == 'adamw': self._compute_adamw(group, info, gather_list, rank, world_size) elif group['kind'] == 'muon': self._compute_muon(group, info, gather_list, rank) else: raise ValueError(f"Unknown optimizer kind: {group['kind']}") # Phase 3: wait for gathers, copy back self._finish_gathers(gather_list) ================================================ FILE: nanochat/report.py ================================================ """ 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) # Return stdout if we got output (even if some files in xargs failed) if result.stdout.strip(): return result.stdout.strip() if result.returncode == 0: return "" 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: count lines/chars in git-tracked source files only extensions = ['py', 'md', 'rs', 'html', 'toml', 'sh'] git_patterns = ' '.join(f"'*.{ext}'" for ext in extensions) files_output = run_command(f"git ls-files -- {git_patterns}") file_list = [f for f in (files_output or '').split('\n') if f] num_files = len(file_list) num_lines = 0 num_chars = 0 if num_files > 0: wc_output = run_command(f"git ls-files -- {git_patterns} | xargs wc -lc 2>/dev/null") if wc_output: total_line = wc_output.strip().split('\n')[-1] parts = total_line.split() if len(parts) >= 2: num_lines = int(parts[0]) num_chars = int(parts[1]) 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", "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-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", "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() ================================================ FILE: nanochat/tokenizer.py ================================================ """ 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 verified that 2 is the sweet spot for vocab size of 32K. 1 is a bit worse, 3 was worse still. 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, num_threads=None): # encode a single string # prepend/append can be either a string of a special token or a token id directly. # num_threads is ignored (only used by the nanochat Tokenizer for parallel encoding) 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): # Different HuggingFace models use different BOS tokens and there is little consistency # 1) attempt to find a <|bos|> token bos = self.encode_special("<|bos|>") # 2) if that fails, attempt to find a <|endoftext|> token (e.g. GPT-2 models) if bos is None: bos = self.encode_special("<|endoftext|>") # 3) if these fail, it's better to crash than to silently return None assert bos is not None, "Failed to find BOS token in tokenizer" 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 ================================================ FILE: nanochat/ui.html ================================================ NanoChat

nanochat

================================================ FILE: pyproject.toml ================================================ [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", "ipykernel>=7.1.0", "kernels>=0.11.7", "matplotlib>=3.10.8", "psutil>=7.1.0", "python-dotenv>=1.2.1", "regex>=2025.9.1", "rustbpe>=0.1.0", "scipy>=1.15.3", "setuptools>=80.9.0", "tabulate>=0.9.0", "tiktoken>=0.11.0", "tokenizers>=0.22.0", "torch==2.9.1", "transformers>=4.57.3", "uvicorn>=0.36.0", "wandb>=0.21.3", "zstandard>=0.25.0", ] [dependency-groups] dev = [ "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.9.1", ] gpu = [ "torch==2.9.1", ] [tool.uv] conflicts = [ [ { extra = "cpu" }, { extra = "gpu" }, ], ] ================================================ FILE: runs/miniseries.sh ================================================ #!/bin/bash # See speedrun.sh for more comments # Usage: ./miniseries.sh [series_name] # Example: ./miniseries.sh jan11 # Default series name is today's date (e.g., jan11) export OMP_NUM_THREADS=1 export NANOCHAT_BASE_DIR="$HOME/.cache/nanochat" mkdir -p $NANOCHAT_BASE_DIR # Setup (skip with SKIP_SETUP=1) if [ -z "$SKIP_SETUP" ]; then # uv 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 # Tokenizer, download 1000 shards for pretraining # (probably this can be reduced but it's tricky to determine the exact right number, TODO). python -m nanochat.dataset -n 1000 python -m scripts.tok_train --max-chars=2000000000 --vocab-size=32768 else source .venv/bin/activate fi # Series name: from arg, env var, or default to today's date (e.g., jan11) SERIES_NAME="${1:-${SERIES_NAME:-$(date +%b%d | tr '[:upper:]' '[:lower:]')}}" # Depths to train (the "miniseries") DEPTHS=(12 14 16 18 20 22 24 26) # Hardware NPROC_PER_NODE="${NPROC_PER_NODE:-8}" # Logging WANDB_RUN="${WANDB_RUN:-${SERIES_NAME}_miniseries}" RESULTS_DIR="$NANOCHAT_BASE_DIR/${SERIES_NAME}_miniseries_results" mkdir -p "$RESULTS_DIR" RESULTS_FILE="$RESULTS_DIR/results.csv" # Write CSV header only if file doesn't exist if [ ! -f "$RESULTS_FILE" ]; then echo "depth,model_dim,num_params,num_scaling_params,num_iterations,tokens_trained,param_data_ratio,val_bpb,core_score,train_time_sec" > "$RESULTS_FILE" fi log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" } log "==============================================" log "${SERIES_NAME} Miniseries Training" log "==============================================" for d in "${DEPTHS[@]}"; do log "Training d=$d..." TAG="${SERIES_NAME}_miniseries_d${d}" START_TIME=$(date +%s) # Reduce --device-batch-size to avoid OOM at larger depths if [ $d -ge 28 ]; then DEVICE_BATCH_SIZE_ARG="--device-batch-size=8" elif [ $d -ge 20 ]; then DEVICE_BATCH_SIZE_ARG="--device-batch-size=16" else DEVICE_BATCH_SIZE_ARG="--device-batch-size=32" fi torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_train -- \ --depth=$d \ --run="${WANDB_RUN}_d${d}" \ --model-tag="${TAG}" \ --core-metric-every=999999 \ --core-metric-max-per-task=-1 \ --sample-every=-1 \ --save-every=-1 \ $DEVICE_BATCH_SIZE_ARG \ 2>&1 | tee "$RESULTS_DIR/${TAG}_train.log" END_TIME=$(date +%s) TRAIN_TIME=$((END_TIME - START_TIME)) # Extract stats from log LOG_FILE="$RESULTS_DIR/${TAG}_train.log" NUM_PARAMS=$(grep "Number of parameters:" "$LOG_FILE" | tail -1 | grep -oP '[\d,]+' | head -1 | tr -d ',') NUM_SCALING_PARAMS=$(grep "Number of parameters:" "$LOG_FILE" | tail -1 | grep -oP 'scaling: [\d,]+' | grep -oP '[\d,]+' | tr -d ',') NUM_ITERS=$(grep "Calculated number of iterations" "$LOG_FILE" | tail -1 | sed 's/.*: //' | tr -d ',') TOKENS_TRAINED=$((NUM_ITERS * 524288)) PARAM_DATA_RATIO=$(python -c "print(f'{$TOKENS_TRAINED / $NUM_SCALING_PARAMS:.2f}')") MODEL_DIM=$((d * 64)) VAL_BPB=$(grep "Validation bpb:" "$LOG_FILE" | tail -1 | grep -oP '[\d.]+$') CORE_SCORE=$(grep "CORE metric:" "$LOG_FILE" | tail -1 | awk '{print $NF}') if [ -z "$CORE_SCORE" ]; then CORE_SCORE="0.0" fi log " d=$d: params=$NUM_PARAMS, scaling=$NUM_SCALING_PARAMS, ratio=$PARAM_DATA_RATIO, bpb=$VAL_BPB, CORE=$CORE_SCORE, time=${TRAIN_TIME}s" # Append to CSV echo "$d,$MODEL_DIM,$NUM_PARAMS,$NUM_SCALING_PARAMS,$NUM_ITERS,$TOKENS_TRAINED,$PARAM_DATA_RATIO,$VAL_BPB,$CORE_SCORE,$TRAIN_TIME" >> "$RESULTS_FILE" done log "==============================================" log "${SERIES_NAME} Miniseries Complete!" log "==============================================" log "Results saved to: $RESULTS_FILE" echo "" echo "Results:" column -t -s',' "$RESULTS_FILE" ================================================ FILE: runs/runcpu.sh ================================================ #!/bin/bash # Showing an example run for exercising some of the code paths on the CPU (or MPS on Macbooks) # This script was last updated/tuned on Jan 17, 2026. # Run as: # bash runs/runcpu.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. # You may also want to run this script manually and one by one, copy pasting commands into your terminal. # all the setup stuff 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 # train tokenizer on ~2B characters (~34 seconds on my MacBook Pro M3 Max) python -m nanochat.dataset -n 8 python -m scripts.tok_train --max-chars=2000000000 python -m scripts.tok_eval # train a small 4 layer model # I tuned this run to complete in about 30 minutes on my MacBook Pro M3 Max. # To get better results, try increasing num_iterations, or get other ideas from your favorite LLM. python -m scripts.base_train \ --depth=6 \ --head-dim=64 \ --window-pattern=L \ --max-seq-len=512 \ --device-batch-size=32 \ --total-batch-size=16384 \ --eval-every=100 \ --eval-tokens=524288 \ --core-metric-every=-1 \ --sample-every=100 \ --num-iterations=5000 \ --run=$WANDB_RUN python -m scripts.base_eval --device-batch-size=1 --split-tokens=16384 --max-per-task=16 # SFT (~10 minutes on my MacBook Pro M3 Max) curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl python -m scripts.chat_sft \ --max-seq-len=512 \ --device-batch-size=32 \ --total-batch-size=16384 \ --eval-every=200 \ --eval-tokens=524288 \ --num-iterations=1500 \ --run=$WANDB_RUN # Chat with the model over CLI # The model should be able to say that it is Paris. # It might even know that the color of the sky is blue. # Sometimes the model likes it if you first say Hi before you ask it questions. # python -m scripts.chat_cli -p "What is the capital of France?" # Chat with the model over a pretty WebUI ChatGPT style # python -m scripts.chat_web ================================================ FILE: runs/scaling_laws.sh ================================================ #!/bin/bash LABEL="jan26" FLOPS_BUDGETS=( 1e18 2.15e18 4.64e18 1e19 ) DEPTHS=(8 10 12 14 16 18 20) NPROC_PER_NODE="${NPROC_PER_NODE:-8}" WANDB_RUN="${WANDB_RUN:-scaling_${LABEL}}" EVAL_TOKENS=$((100 * 524288)) # ~100M tokens for final eval (default is ~10M) export OMP_NUM_THREADS=1 export NANOCHAT_BASE_DIR="${NANOCHAT_BASE_DIR:-$HOME/.cache/nanochat}" source .venv/bin/activate RESULTS_DIR="$NANOCHAT_BASE_DIR/scaling_laws_results_${LABEL}" mkdir -p "$RESULTS_DIR" RESULTS_FILE="$RESULTS_DIR/results.csv" # Write CSV header only if file doesn't exist if [ ! -f "$RESULTS_FILE" ]; then echo "flops_budget,depth,model_dim,params_wte,params_value_embeds,params_lm_head,params_transformer,params_scalars,params_total,num_iterations,tokens_trained,val_bpb,core_score,train_time_sec" > "$RESULTS_FILE" fi log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" } # Check if a run already exists in results run_exists() { local flops=$1 local depth=$2 grep -q "^${flops},${depth}," "$RESULTS_FILE" 2>/dev/null } # ============================================================================= # Main Loop # ============================================================================= for flops in "${FLOPS_BUDGETS[@]}"; do log "==============================================" log "Compute budget: $flops FLOPs" log "==============================================" for d in "${DEPTHS[@]}"; do # Skip if already completed if run_exists "$flops" "$d"; then log "Skipping d=$d at $flops FLOPs (already in results)" continue fi log "Training d=$d at $flops FLOPs..." # Unique tag for this run TAG="scaling_${flops}_d${d}" # Record start time START_TIME=$(date +%s) # Train the model with fixed flops budget # The script will auto-calculate num_iterations to hit target_flops # CORE eval happens once at the end (999999 ensures only final step) torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_train -- \ --depth=$d \ --target-flops=$flops \ --target-param-data-ratio=-1 \ --run="${WANDB_RUN}_${TAG}" \ --model-tag="${TAG}" \ --eval-tokens=$EVAL_TOKENS \ --core-metric-every=999999 \ --core-metric-max-per-task=-1 \ --sample-every=-1 \ --save-every=-1 \ 2>&1 | tee "$RESULTS_DIR/${TAG}_train.log" END_TIME=$(date +%s) TRAIN_TIME=$((END_TIME - START_TIME)) # Extract training stats from the log LOG_FILE="$RESULTS_DIR/${TAG}_train.log" # Extract detailed parameter counts (for scaling law analysis with different conventions) # Note: the log format is padded, e.g. "wte : 25,165,824" # so we grep for "^key " (key at start of line followed by space) to avoid false matches PARAMS_WTE=$(grep "^wte " "$LOG_FILE" | tail -1 | grep -oP '[\d,]+' | tr -d ',') PARAMS_VE=$(grep "^value_embeds " "$LOG_FILE" | tail -1 | grep -oP '[\d,]+' | tr -d ',') PARAMS_LM=$(grep "^lm_head " "$LOG_FILE" | tail -1 | grep -oP '[\d,]+' | tr -d ',') PARAMS_TRANSFORMER=$(grep "^transformer_matrices " "$LOG_FILE" | tail -1 | grep -oP '[\d,]+' | tr -d ',') PARAMS_SCALARS=$(grep "^scalars " "$LOG_FILE" | tail -1 | grep -oP '[\d,]+' | tr -d ',') PARAMS_TOTAL=$(grep "^total " "$LOG_FILE" | tail -1 | grep -oP '[\d,]+' | tr -d ',') NUM_ITERS=$(grep "Calculated number of iterations" "$LOG_FILE" | tail -1 | sed 's/.*: //' | tr -d ',') # Calculate tokens trained (iterations * batch_size, default 524288) TOKENS_TRAINED=$((NUM_ITERS * 524288)) # Model dim MODEL_DIM=$((d * 64)) # Val BPB from final eval VAL_BPB=$(grep "Validation bpb:" "$LOG_FILE" | tail -1 | grep -oP '[\d.]+$') # Extract CORE score from training log (evaluated on final step) CORE_SCORE=$(grep "CORE metric:" "$LOG_FILE" | tail -1 | awk '{print $NF}') if [ -z "$CORE_SCORE" ]; then log "WARNING: Could not extract CORE score for d=$d" CORE_SCORE="0.0" fi log " Params: $PARAMS_TOTAL (transformer: $PARAMS_TRANSFORMER), Iters: $NUM_ITERS, Val BPB: $VAL_BPB, CORE: $CORE_SCORE" # Append to CSV echo "$flops,$d,$MODEL_DIM,$PARAMS_WTE,$PARAMS_VE,$PARAMS_LM,$PARAMS_TRANSFORMER,$PARAMS_SCALARS,$PARAMS_TOTAL,$NUM_ITERS,$TOKENS_TRAINED,$VAL_BPB,$CORE_SCORE,$TRAIN_TIME" >> "$RESULTS_FILE" done done log "==============================================" log "Scaling Laws Sweep Complete" log "==============================================" log "Results saved to: $RESULTS_FILE" echo "" echo "Results:" column -t -s',' "$RESULTS_FILE" ================================================ FILE: runs/speedrun.sh ================================================ #!/bin/bash # This script is configured to train your own GPT-2 grade LLM (pretraining + finetuning) # It is designed to run on a blank 8XH100 GPU node and takes approximately 3 hours to complete. # 1) Example launch (simplest): # bash runs/speedrun.sh # 2) Example launch in a screen session (because the run takes ~3 hours): # screen -L -Logfile runs/speedrun.log -S speedrun bash runs/speedrun.sh # 3) Example launch with wandb logging, but see below for setting up wandb first: # WANDB_RUN=speedrun screen -L -Logfile runs/speedrun.log -S speedrun bash runs/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 # Download the first ~2B characters of pretraining dataset # 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 # look at dev/repackage_data_reference.py for details on how this data was prepared python -m nanochat.dataset -n 8 # Immediately also kick off downloading more shards in the background while tokenizer trains # Approximately 150 shards are needed for GPT-2 capability pretraining, add 20 for padding. # The maximum total number of shards available in the entire dataset is 6542. python -m nanochat.dataset -n 170 & DATASET_DOWNLOAD_PID=$! # train the tokenizer with vocab size 2**15 = 32768 on ~2B characters of data python -m scripts.tok_train # evaluate the tokenizer (report compression ratio etc.) python -m scripts.tok_eval # ----------------------------------------------------------------------------- # Base model (pretraining) echo "Waiting for dataset download to complete..." wait $DATASET_DOWNLOAD_PID # d24 model (slightly undertrained to beat GPT-2 => decrease data:params ratio from compute optimal 10.5 (default) to 8) torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- --depth=24 --target-param-data-ratio=8 --device-batch-size=16 --fp8 --run=$WANDB_RUN # evaluate the model: CORE metric, BPB on train/val, and draw samples torchrun --standalone --nproc_per_node=8 -m scripts.base_eval -- --device-batch-size=16 # ----------------------------------------------------------------------------- # SFT (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_synthetic_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 SFT and eval the model torchrun --standalone --nproc_per_node=8 -m scripts.chat_sft -- --device-batch-size=16 --run=$WANDB_RUN torchrun --standalone --nproc_per_node=8 -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 # ----------------------------------------------------------------------------- # 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 ================================================ FILE: scripts/base_eval.py ================================================ """ Unified evaluation script for base models. Supports three evaluation modes (comma-separated): --eval core : CORE metric (accuracy on ICL tasks) --eval bpb : Bits per byte on train/val splits --eval sample : Generate samples from the model Default is all three: --eval core,bpb,sample Examples: # Evaluate a HuggingFace model (e.g. GPT-2 124M) using 8 GPUs torchrun --nproc_per_node=8 -m scripts.base_eval --hf-path openai-community/gpt2 # Evaluate a nanochat model (e.g. d24) using 8 GPUs torchrun --nproc_per_node=8 -m scripts.base_eval --model-tag d24 --device-batch-size=16 # Quick/approximate evaluation using a single GPU python -m scripts.base_eval --model-tag d24 --device-batch-size=16 --max-per-task=100 --split-tokens=524288 """ import os import csv import time import json import yaml import shutil import random import zipfile import tempfile import argparse 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, get_token_bytes from nanochat.checkpoint_manager import load_model from nanochat.core_eval import evaluate_task from nanochat.dataloader import tokenizing_distributed_data_loader_bos_bestfit from nanochat.loss_eval import evaluate_bpb from nanochat.engine import Engine # ----------------------------------------------------------------------------- # HuggingFace loading utilities class ModelWrapper: """Lightweight wrapper to give HuggingFace models a nanochat-compatible interface.""" def __init__(self, model, max_seq_len=None): self.model = model self.max_seq_len = max_seq_len def __call__(self, input_ids, targets=None, loss_reduction='mean'): logits = self.model(input_ids).logits if targets is None: return logits loss = torch.nn.functional.cross_entropy( logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction ) return loss def get_device(self): return next(self.model.parameters()).device def load_hf_model(hf_path: str, device): """Load a HuggingFace model and tokenizer.""" print0(f"Loading HuggingFace model from: {hf_path}") from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained(hf_path) model.to(device) model.eval() max_seq_len = 1024 if "gpt2" in hf_path else None model = ModelWrapper(model, max_seq_len=max_seq_len) tokenizer = HuggingFaceTokenizer.from_pretrained(hf_path) return model, tokenizer def get_hf_token_bytes(tokenizer, device="cpu"): """Compute token_bytes tensor for a HuggingFace tokenizer.""" vocab_size = tokenizer.tokenizer.get_vocab_size() token_bytes = torch.zeros(vocab_size, dtype=torch.int64, device=device) for token_id in range(vocab_size): token_str = tokenizer.tokenizer.decode([token_id]) token_bytes[token_id] = len(token_str.encode('utf-8')) return token_bytes # ----------------------------------------------------------------------------- # CORE evaluation EVAL_BUNDLE_URL = "https://karpathy-public.s3.us-west-2.amazonaws.com/eval_bundle.zip" def place_eval_bundle(file_path): """Unzip eval_bundle.zip 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_core(model, tokenizer, device, max_per_task=-1): """ Evaluate a base model on the CORE benchmark. Returns dict with results, centered_results, and core_metric. """ base_dir = get_base_dir() eval_bundle_dir = os.path.join(base_dir, "eval_bundle") # Download the eval bundle 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 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='') 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 for consistent subsampling when using max_per_task shuffle_rng = random.Random(1337) shuffle_rng.shuffle(data) if max_per_task > 0: data = data[:max_per_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 elapsed = time.time() - start_time print0(f"accuracy: {accuracy:.4f} | centered: {centered_result:.4f} | time: {elapsed:.2f}s") core_metric = sum(centered_results.values()) / len(centered_results) out = { "results": results, "centered_results": centered_results, "core_metric": core_metric } return out # ----------------------------------------------------------------------------- # Main def main(): parser = argparse.ArgumentParser(description="Base model evaluation") parser.add_argument('--eval', type=str, default='core,bpb,sample', help='Comma-separated evaluations to run: core,bpb,sample (default: all)') parser.add_argument('--hf-path', type=str, default=None, help='HuggingFace model path (e.g. openai-community/gpt2-xl)') parser.add_argument('--model-tag', type=str, default=None, help='nanochat model tag to identify the checkpoint directory') parser.add_argument('--step', type=int, default=None, help='Model step to load (default = last)') parser.add_argument('--max-per-task', type=int, default=-1, help='Max examples per CORE task (-1 = all)') parser.add_argument('--device-batch-size', type=int, default=32, help='Per-device batch size for BPB evaluation') parser.add_argument('--split-tokens', type=int, default=40*524288, help='Number of tokens to evaluate per split for BPB') parser.add_argument('--device-type', type=str, default='', help='cuda|cpu|mps (empty = autodetect)') args = parser.parse_args() # Parse evaluation modes eval_modes = set(mode.strip() for mode in args.eval.split(',')) valid_modes = {'core', 'bpb', 'sample'} invalid = eval_modes - valid_modes if invalid: parser.error(f"Invalid eval modes: {invalid}. Valid: {valid_modes}") # Distributed / precision setup 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) # Load model and tokenizer is_hf_model = args.hf_path is not None if is_hf_model: model, tokenizer = load_hf_model(args.hf_path, device) sequence_len = model.max_seq_len or 1024 token_bytes = get_hf_token_bytes(tokenizer, device=device) model_name = args.hf_path model_slug = args.hf_path.replace("/", "-") else: model, tokenizer, meta = load_model("base", device, phase="eval", model_tag=args.model_tag, step=args.step) sequence_len = meta["model_config"]["sequence_len"] token_bytes = get_token_bytes(device=device) model_name = f"base_model (step {meta['step']})" model_slug = f"base_model_{meta['step']:06d}" print0(f"Evaluating model: {model_name}") print0(f"Eval modes: {', '.join(sorted(eval_modes))}") # Results to log core_results = None bpb_results = {} samples = [] unconditioned_samples = [] # --- Sampling --- if 'sample' in eval_modes and not is_hf_model: print0("\n" + "="*80) print0("Model Samples") print0("="*80) 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) print0("\nConditioned samples:") for prompt in prompts: tokens = tokenizer(prompt, prepend="<|bos|>") sample, _ = engine.generate_batch(tokens, num_samples=1, max_tokens=16, temperature=0) sample_str = tokenizer.decode(sample[0]) print0("-" * 80) print0(sample_str) samples.append(sample_str) print0("\nUnconditioned samples:") tokens = tokenizer("", prepend="<|bos|>") uncond, _ = engine.generate_batch(tokens, num_samples=8, max_tokens=128, temperature=1.0) for sample in uncond: sample_str = tokenizer.decode(sample) print0("-" * 80) print0(sample_str) unconditioned_samples.append(sample_str) elif 'sample' in eval_modes and is_hf_model: print0("\nSkipping sampling for HuggingFace models (not supported)") # --- BPB evaluation --- if 'bpb' in eval_modes: print0("\n" + "="*80) print0("BPB Evaluation") print0("="*80) tokens_per_step = args.device_batch_size * sequence_len * ddp_world_size if args.split_tokens % tokens_per_step != 0: # Adjust to nearest multiple args.split_tokens = (args.split_tokens // tokens_per_step) * tokens_per_step print0(f"Adjusted split_tokens to {args.split_tokens} (must be divisible by {tokens_per_step})") steps = args.split_tokens // tokens_per_step for split_name in ["train", "val"]: loader = tokenizing_distributed_data_loader_bos_bestfit(tokenizer, args.device_batch_size, sequence_len, split_name, device=device) bpb = evaluate_bpb(model, loader, steps, token_bytes) bpb_results[split_name] = bpb print0(f"{split_name} bpb: {bpb:.6f}") # --- CORE evaluation --- if 'core' in eval_modes: print0("\n" + "="*80) print0("CORE Evaluation") print0("="*80) core_results = evaluate_core(model, tokenizer, device, max_per_task=args.max_per_task) # Write CSV output 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) 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 core_results["results"]: acc = core_results["results"][label] centered = core_results["centered_results"][label] f.write(f"{label:<35}, {acc:<10.6f}, {centered:<10.6f}\n") f.write(f"{'CORE':<35}, {'':<10}, {core_results['core_metric']:<10.6f}\n") print0(f"\nResults written to: {output_csv_path}") print0(f"CORE metric: {core_results['core_metric']:.4f}") # --- Log to report --- from nanochat.report import get_report report_data = [{"model": model_name}] if core_results: report_data[0]["CORE metric"] = core_results["core_metric"] report_data.append(core_results["centered_results"]) if bpb_results: report_data[0]["train bpb"] = bpb_results.get("train") report_data[0]["val bpb"] = bpb_results.get("val") if samples: report_data.append({f"sample {i}": s for i, s in enumerate(samples)}) if unconditioned_samples: report_data.append({f"unconditioned {i}": s for i, s in enumerate(unconditioned_samples)}) get_report().log(section="Base model evaluation", data=report_data) compute_cleanup() if __name__ == "__main__": main() ================================================ FILE: scripts/base_train.py ================================================ """ Train model. From root directory of the project, run as: python -m scripts.base_train or distributed as: torchrun --nproc_per_node=8 -m scripts.base_train 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_ALLOC_CONF"] = "expandable_segments:True" import gc import json import time import math import argparse from dataclasses import asdict from contextlib import contextmanager import wandb import torch import torch.distributed as dist from nanochat.gpt import GPT, GPTConfig, Linear from nanochat.dataloader import tokenizing_distributed_data_loader_bos_bestfit, tokenizing_distributed_data_loader_with_state_bos_bestfit from nanochat.common import compute_init, compute_cleanup, print0, DummyWandb, print_banner, get_base_dir, autodetect_device_type, get_peak_flops, COMPUTE_DTYPE, COMPUTE_DTYPE_REASON, is_ddp_initialized from nanochat.tokenizer import get_tokenizer, get_token_bytes from nanochat.checkpoint_manager import save_checkpoint, load_checkpoint from nanochat.loss_eval import evaluate_bpb from nanochat.engine import Engine from nanochat.flash_attention import HAS_FA3 from scripts.base_eval import evaluate_core print_banner() # ----------------------------------------------------------------------------- # CLI arguments parser = argparse.ArgumentParser(description="Pretrain base model") # Logging parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('dummy' disables wandb logging)") # Runtime parser.add_argument("--device-type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") # FP8 training parser.add_argument("--fp8", action="store_true", help="enable FP8 training (requires H100+ GPU and torchao)") parser.add_argument("--fp8-recipe", type=str, default="tensorwise", choices=["rowwise", "tensorwise"], help="FP8 scaling recipe: tensorwise (faster, recommended) or rowwise (more accurate but slower)") # Model architecture parser.add_argument("--depth", type=int, default=20, help="depth of the Transformer model") parser.add_argument("--aspect-ratio", type=int, default=64, help="model_dim = depth * aspect_ratio") parser.add_argument("--head-dim", type=int, default=128, help="target head dimension for attention") parser.add_argument("--max-seq-len", type=int, default=2048, help="max context length") parser.add_argument("--window-pattern", type=str, default="SSSL", help="sliding window pattern tiled across layers: L=full, S=half context (e.g. 'SSL')") # Training horizon (only one used, in order of precedence) parser.add_argument("--num-iterations", type=int, default=-1, help="explicit number of optimization steps (-1 = disable)") parser.add_argument("--target-flops", type=float, default=-1.0, help="calculate num_iterations to reach target_flops (-1 = disable)") parser.add_argument("--target-param-data-ratio", type=float, default=10.5, help="calculate num_iterations to maintain data:param ratio (Chinchilla=20, -1 = disable)") # Optimization parser.add_argument("--device-batch-size", type=int, default=32, help="per-device batch size. good number to reduce to 16,8,4,... if you OOM on VRAM.") parser.add_argument("--total-batch-size", type=int, default=-1, help="total batch size in tokens. decent numbers are e.g. 524288. (-1 = auto-compute optimal)") parser.add_argument("--embedding-lr", type=float, default=0.3, help="learning rate for embedding parameters (Adam)") parser.add_argument("--unembedding-lr", type=float, default=0.008, help="learning rate for unembedding parameters (Adam)") parser.add_argument("--weight-decay", type=float, default=0.28, help="cautious weight decay for the Muon optimizer (for weights)") parser.add_argument("--matrix-lr", type=float, default=0.02, help="learning rate for matrix parameters (Muon)") parser.add_argument("--scalar-lr", type=float, default=0.5, help="learning rate for scalars (resid_lambdas, x0_lambdas)") parser.add_argument("--warmup-steps", type=int, default=40, help="number of steps for LR warmup") parser.add_argument("--warmdown-ratio", type=float, default=0.65, help="ratio of iterations for LR warmdown") parser.add_argument("--final-lr-frac", type=float, default=0.05, help="final LR as fraction of initial LR") parser.add_argument("--resume-from-step", type=int, default=-1, help="resume training from this step (-1 = disable)") # Evaluation parser.add_argument("--eval-every", type=int, default=250, help="evaluate val bpb every N steps (-1 = disable)") parser.add_argument("--eval-tokens", type=int, default=80*524288, help="number of tokens to evaluate val loss on") parser.add_argument("--core-metric-every", type=int, default=2000, help="evaluate CORE metric every N steps (-1 = disable)") parser.add_argument("--core-metric-max-per-task", type=int, default=500, help="examples per task for CORE metric") parser.add_argument("--sample-every", type=int, default=2000, help="sample from model every N steps (-1 = disable)") parser.add_argument("--save-every", type=int, default=-1, help="save checkpoints every N steps (-1 = only at end)") # Output parser.add_argument("--model-tag", type=str, default=None, help="override model tag for checkpoint directory name") args = parser.parse_args() user_config = vars(args).copy() # for logging # ----------------------------------------------------------------------------- # Compute init and wandb logging 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) master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. 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 if device_type == "cuda": gpu_device_name = torch.cuda.get_device_name(0) gpu_peak_flops = get_peak_flops(gpu_device_name) print0(f"GPU: {gpu_device_name} | Peak FLOPS (BF16): {gpu_peak_flops:.2e}") else: gpu_peak_flops = float('inf') # MFU not meaningful for CPU/MPS print0(f"COMPUTE_DTYPE: {COMPUTE_DTYPE} ({COMPUTE_DTYPE_REASON})") # wandb logging init use_dummy_wandb = args.run == "dummy" or not master_process wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat", name=args.run, config=user_config) # Flash Attention status from nanochat.flash_attention import USE_FA3 using_fa3 = USE_FA3 if using_fa3: print0("✓ Using Flash Attention 3 (Hopper GPU detected), efficient, new and awesome.") else: print0("!" * 80) if HAS_FA3 and COMPUTE_DTYPE != torch.bfloat16: print0(f"WARNING: Flash Attention 3 only supports bf16, but COMPUTE_DTYPE={COMPUTE_DTYPE}. Using PyTorch SDPA fallback") else: print0("WARNING: Flash Attention 3 not available, using PyTorch SDPA fallback") print0("WARNING: Training will be less efficient without FA3") if args.window_pattern != "L": print0(f"WARNING: SDPA has no support for sliding window attention (window_pattern='{args.window_pattern}'). Your GPU utilization will be terrible.") print0("WARNING: Recommend using --window-pattern L for full context attention without alternating sliding window patterns.") print0("!" * 80) # ----------------------------------------------------------------------------- # Tokenizer will be useful for evaluation and also we need the vocab size to init the model tokenizer = get_tokenizer() token_bytes = get_token_bytes(device=device) vocab_size = tokenizer.get_vocab_size() print0(f"Vocab size: {vocab_size:,}") # ----------------------------------------------------------------------------- # Initialize the Model def build_model_meta(depth): """Build a model on meta device for a given depth (shapes/dtypes only, no data).""" # Model dim is nudged up to nearest multiple of head_dim for clean division # (FA3 requires head_dim divisible by 8, and this guarantees head_dim == args.head_dim exactly) base_dim = depth * args.aspect_ratio model_dim = ((base_dim + args.head_dim - 1) // args.head_dim) * args.head_dim num_heads = model_dim // args.head_dim config = GPTConfig( sequence_len=args.max_seq_len, vocab_size=vocab_size, n_layer=depth, n_head=num_heads, n_kv_head=num_heads, n_embd=model_dim, window_pattern=args.window_pattern, ) with torch.device("meta"): model_meta = GPT(config) return model_meta # Build the model, move to device, init the weights model = build_model_meta(args.depth) # 1) Build on meta device (only shapes/dtypes, no data) model_config = model.config model_config_kwargs = asdict(model_config) print0(f"Model config:\n{json.dumps(model_config_kwargs, indent=2)}") model.to_empty(device=device) # 2) All tensors get storage on target device but with uninitialized (garbage) data model.init_weights() # 3) All tensors get initialized # If we are resuming, overwrite the model parameters with those of the checkpoint base_dir = get_base_dir() output_dirname = args.model_tag if args.model_tag else f"d{args.depth}" # e.g. d12 checkpoint_dir = os.path.join(base_dir, "base_checkpoints", output_dirname) resuming = args.resume_from_step != -1 if resuming: print0(f"Resuming optimization from step {args.resume_from_step}") model_data, optimizer_data, meta_data = load_checkpoint(checkpoint_dir, args.resume_from_step, device, load_optimizer=True, rank=ddp_rank) model.load_state_dict(model_data, strict=True, assign=True) del model_data # free up this memory after the copy # ----------------------------------------------------------------------------- # FP8 training initialization and management (this has to be done before torch.compile) # Convert Linear layers to Float8Linear if --fp8 is set if args.fp8: if device_type != "cuda": print0("Warning: FP8 training requires CUDA, ignoring --fp8 flag") else: # our custom fp8 is simpler than torchao, written for exact API compatibility from nanochat.fp8 import Float8LinearConfig, convert_to_float8_training # from torchao.float8 import Float8LinearConfig, convert_to_float8_training import torch.nn as nn # Filter: dims must be divisible by 16 (FP8 hardware requirement) large enough def fp8_module_filter(mod: nn.Module, fqn: str) -> bool: if not isinstance(mod, nn.Linear): return False if mod.in_features % 16 != 0 or mod.out_features % 16 != 0: return False if min(mod.in_features, mod.out_features) < 128: return False return True fp8_config = Float8LinearConfig.from_recipe_name(args.fp8_recipe) num_linear = sum(1 for m in model.modules() if isinstance(m, nn.Linear)) convert_to_float8_training(model, config=fp8_config, module_filter_fn=fp8_module_filter) num_fp8 = sum(1 for m in model.modules() if 'Float8' in type(m).__name__) num_skipped = num_linear - num_fp8 print0(f"✓ FP8 training enabled ({args.fp8_recipe} scaling) - converted {num_fp8}/{num_linear} linear layers, skipped {num_skipped} (too small)") # Context manager to temporarily disable FP8 so that model evaluation remains in BF16 @contextmanager def disable_fp8(model): """Temporarily swap Float8Linear modules with nn.Linear for BF16 evaluation. CastConfig is a frozen dataclass, so we can't mutate scaling_type. Instead, we swap out Float8Linear modules entirely and restore them after. """ import torch.nn as nn # Find all Float8Linear modules and their locations fp8_locations = [] # list of (parent_module, attr_name, fp8_module) for name, module in model.named_modules(): if 'Float8' in type(module).__name__: if '.' in name: parent_name, attr_name = name.rsplit('.', 1) parent = model.get_submodule(parent_name) else: parent = model attr_name = name fp8_locations.append((parent, attr_name, module)) if not fp8_locations: yield # No FP8 modules, nothing to do return # Swap Float8Linear -> Linear (our custom class that casts weights to match input dtype) for parent, attr_name, fp8_module in fp8_locations: linear = Linear( fp8_module.in_features, fp8_module.out_features, bias=fp8_module.bias is not None, device=fp8_module.weight.device, dtype=fp8_module.weight.dtype, ) linear.weight = fp8_module.weight # share, don't copy if fp8_module.bias is not None: linear.bias = fp8_module.bias setattr(parent, attr_name, linear) try: yield finally: # Restore Float8Linear modules for parent, attr_name, fp8_module in fp8_locations: setattr(parent, attr_name, fp8_module) # ----------------------------------------------------------------------------- # Compile the model orig_model = model # original, uncompiled model, for saving raw model state_dict and for inference/evaluation (because the shapes may change shape) model = torch.compile(model, dynamic=False) # the inputs to model will never change shape so dynamic=False is safe # ----------------------------------------------------------------------------- # Scaling laws and muP extrapolations to determine the optimal training horizon, batch size, learning rates, weight decay. # Get the parameter counts of our model param_counts = model.num_scaling_params() print0(f"Parameter counts:") for key, value in param_counts.items(): print0(f"{key:24s}: {value:,}") num_params = param_counts['total'] num_flops_per_token = model.estimate_flops() print0(f"Estimated FLOPs per token: {num_flops_per_token:e}") # 1) Use scaling laws to determine the optimal training horizon in tokens # The compute-optimal models satisfy the Tokens:Params ratio of --target-param-data-ratio (derived experimentally via scaling laws analysis). # We've already initialized the model so we have Params. Optimal Tokens is now simply target-param-data-ratio * Params def get_scaling_params(m): # As for which params to use exactly, transformer matrices + lm_head gives cleanest scaling laws (see dev/LOG.md Jan 27, 2026) params_counts = m.num_scaling_params() scaling_params = params_counts['transformer_matrices'] + params_counts['lm_head'] return scaling_params num_scaling_params = get_scaling_params(model) target_tokens = int(args.target_param_data_ratio * num_scaling_params) # optimal tokens for the model we are about to train # Our reference model is d12, this is where a lot of hyperparameters are tuned and then transfered to higher depths (muP style) d12_ref = build_model_meta(12) # creates the model on meta device D_REF = args.target_param_data_ratio * get_scaling_params(d12_ref) # compute-optimal d12 training horizon in tokens (measured empirically) B_REF = 2**19 # optimal batch size at d12 ~= 524,288 tokens (measured empirically) # 2) Now that we have the token horizon, we can calculate the optimal batch size # We follow the Power Lines paper (Bopt ∝ D^0.383), ref: https://arxiv.org/abs/2505.13738 # The optimal batch size grows as approximately D^0.383, so e.g. if D doubles from d12 to d24, B should grow by 2^0.383 ≈ 1.3x. total_batch_size = args.total_batch_size # user-provided override is possible if total_batch_size == -1: batch_size_ratio = target_tokens / D_REF predicted_batch_size = B_REF * batch_size_ratio ** 0.383 total_batch_size = 2 ** round(math.log2(predicted_batch_size)) # clamp to nearest power of 2 for efficiency print0(f"Auto-computed optimal batch size: {total_batch_size:,} tokens") # 3) Knowing the batch size, we can now calculate a learning rate correction (bigger batch size allows higher learning rates) batch_lr_scale = 1.0 batch_ratio = total_batch_size / B_REF # B/B_ref if batch_ratio != 1.0: # SGD: linear scaling with batch size is standard (not used in nanochat) # AdamW: sqrt scaling is standard: η ∝ √(B/B_ref) # Muon: we will use the same scaling for Muon as for AdamW: η ∝ √(B/B_ref) (not studied carefully, assumption!) batch_lr_scale = batch_ratio ** 0.5 # η ∝ √(B/B_ref) print0(f"Scaling LRs by {batch_lr_scale:.4f} for batch size {total_batch_size:,} (reference: {B_REF:,})") # 4) Knowing the batch size and the token horizon, we can now calculate the appropriate weight decay scaling # We adopt the T_epoch framework from https://arxiv.org/abs/2405.13698 # Central idea of the paper is that T_epoch = B/(η·λ·D) should remain constant. # Above, we used learning rate scaling η ∝ √(B/B_ref). So it's a matter of ~10 lines of math to derive that to keep T_epoch constant, we need: # λ = λ_ref · √(B/B_ref) · (D_ref/D) # Note that these papers study AdamW, *not* Muon. We are blindly following AdamW theory for scaling hoping it ~works for Muon too. weight_decay_scaled = args.weight_decay * math.sqrt(total_batch_size / B_REF) * (D_REF / target_tokens) if weight_decay_scaled != args.weight_decay: print0(f"Scaling weight decay from {args.weight_decay:.6f} to {weight_decay_scaled:.6f} for depth {args.depth}") # ----------------------------------------------------------------------------- # Initialize the Optimizer (combined MuonAdamW: Muon for matrix params, AdamW for rest) optimizer = model.setup_optimizer( # AdamW hyperparameters unembedding_lr=args.unembedding_lr * batch_lr_scale, embedding_lr=args.embedding_lr * batch_lr_scale, scalar_lr=args.scalar_lr * batch_lr_scale, # Muon hyperparameters matrix_lr=args.matrix_lr * batch_lr_scale, weight_decay=weight_decay_scaled, ) if resuming: optimizer.load_state_dict(optimizer_data) del optimizer_data # ----------------------------------------------------------------------------- # GradScaler for fp16 training (bf16/fp32 don't need it — bf16 has the same exponent range as fp32) scaler = torch.amp.GradScaler() if COMPUTE_DTYPE == torch.float16 else None if scaler is not None: print0("GradScaler enabled for fp16 training") # ----------------------------------------------------------------------------- # Initialize the DataLoaders for train/val dataloader_resume_state_dict = None if not resuming else meta_data["dataloader_state_dict"] train_loader = tokenizing_distributed_data_loader_with_state_bos_bestfit(tokenizer, args.device_batch_size, args.max_seq_len, split="train", device=device, resume_state_dict=dataloader_resume_state_dict) build_val_loader = lambda: tokenizing_distributed_data_loader_bos_bestfit(tokenizer, args.device_batch_size, args.max_seq_len, split="val", device=device) x, y, dataloader_state_dict = next(train_loader) # kick off load of the very first batch of data # ----------------------------------------------------------------------------- # Calculate the number of iterations we will train for and set up the various schedulers # num_iterations: either it is given, or from target flops, or from target data:param ratio (in that order) assert args.num_iterations > 0 or args.target_param_data_ratio > 0 or args.target_flops > 0 if args.num_iterations > 0: # Override num_iterations to a specific value if given num_iterations = args.num_iterations print0(f"Using user-provided number of iterations: {num_iterations:,}") elif args.target_flops > 0: # Calculate the number of iterations from the target flops (used in scaling laws analysis, e.g. runs/scaling_laws.sh) num_iterations = round(args.target_flops / (num_flops_per_token * total_batch_size)) print0(f"Calculated number of iterations from target FLOPs: {num_iterations:,}") elif args.target_param_data_ratio > 0: # Calculate the number of iterations from the target param data ratio (the most common use case) 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 # the actual number of tokens we will train for print0(f"Total number of training tokens: {total_tokens:,}") print0(f"Tokens : Scaling params ratio: {total_batch_size * num_iterations / num_scaling_params:.2f}") # e.g. Chinchilla was ~20 print0(f"Total training FLOPs estimate: {num_flops_per_token * total_tokens:e}") # Learning rate schedule (linear warmup, constant, linear warmdown) def get_lr_multiplier(it): warmup_iters = args.warmup_steps warmdown_iters = round(args.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) * args.final_lr_frac # Momentum scheduler for Muon optimizer (warms up to 0.97, warms down to 0.90 during LR warmdown) def get_muon_momentum(it): warmdown_iters = round(args.warmdown_ratio * num_iterations) warmdown_start = num_iterations - warmdown_iters if it < 400: frac = it / 400 return (1 - frac) * 0.85 + frac * 0.97 elif it >= warmdown_start: progress = (it - warmdown_start) / warmdown_iters return 0.97 * (1 - progress) + 0.90 * progress else: return 0.97 # Weight decay scheduler for Muon optimizer (cosine decay to zero over the course of training) def get_weight_decay(it): return weight_decay_scaled * 0.5 * (1 + math.cos(math.pi * it / num_iterations)) # ----------------------------------------------------------------------------- # Training loop # Loop state (variables updated by the training loop) if not resuming: step = 0 val_bpb = None # will be set if eval_every > 0 min_val_bpb = float("inf") smooth_train_loss = 0 # EMA of training loss total_training_time = 0 # total wall-clock time of training else: step = meta_data["step"] loop_state = meta_data["loop_state"] val_bpb = meta_data["val_bpb"] min_val_bpb = loop_state["min_val_bpb"] smooth_train_loss = loop_state["smooth_train_loss"] total_training_time = loop_state["total_training_time"] # Figure out the needed gradient accumulation micro-steps to reach the desired total batch size per step tokens_per_fwdbwd = args.device_batch_size * args.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: {args.device_batch_size} x {args.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}") # Go! while True: last_step = step == num_iterations # loop runs num_iterations+1 times so that we can eval/save at the end flops_so_far = num_flops_per_token * total_batch_size * step # once in a while: evaluate the val bpb (all ranks participate) if args.eval_every > 0 and (last_step or step % args.eval_every == 0): model.eval() val_loader = build_val_loader() eval_steps = args.eval_tokens // (args.device_batch_size * args.max_seq_len * ddp_world_size) with disable_fp8(model): val_bpb = evaluate_bpb(model, val_loader, eval_steps, token_bytes) print0(f"Step {step:05d} | Validation bpb: {val_bpb:.6f}") 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 # disable FP8 for evaluation to use BF16 for more consistent/accurate results results = {} if args.core_metric_every > 0 and (last_step or (step > 0 and step % args.core_metric_every == 0)): model.eval() with disable_fp8(orig_model): results = evaluate_core(orig_model, tokenizer, device, max_per_task=args.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 args.sample_every > 0 and master_process and (last_step or (step > 0 and step % args.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 disable_fp8(orig_model): 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, or every save_every steps, except at the first step or the resume step if last_step or (step > 0 and step != args.resume_from_step and args.save_every > 0 and step % args.save_every == 0): save_checkpoint( checkpoint_dir, step, orig_model.state_dict(), # model parameters optimizer.state_dict(), # optimizer state { # metadata saved as json "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": args.device_batch_size, "max_seq_len": args.max_seq_len, "total_batch_size": total_batch_size, "dataloader_state_dict": dataloader_state_dict, "loop_state": { # all loop state (other than step) so that we can resume training "min_val_bpb": min_val_bpb, "smooth_train_loss": smooth_train_loss, "total_training_time": total_training_time, }, }, rank=ddp_rank, ) # termination conditions (TODO: possibly also add loss explosions etc.) if last_step: break # ------------------------------------------------------------------------- # single training step # evaluate the gradient synchronize() t0 = time.time() for micro_step in range(grad_accum_steps): loss = model(x, y) train_loss = loss.detach() # for logging loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here if scaler is not None: scaler.scale(loss).backward() else: loss.backward() x, y, dataloader_state_dict = next(train_loader) # prefetch the next batch while the GPU is busy with forward/backward # step the optimizer lrm = get_lr_multiplier(step) muon_momentum = get_muon_momentum(step) muon_weight_decay = get_weight_decay(step) for group in optimizer.param_groups: group["lr"] = group["initial_lr"] * lrm if group['kind'] == 'muon': group["momentum"] = muon_momentum group["weight_decay"] = muon_weight_decay if scaler is not None: scaler.unscale_(optimizer) # In distributed training, all ranks must agree on whether to skip the step. # Each rank may independently encounter inf/nan gradients, so we all-reduce # the found_inf flag (MAX = if any rank found inf, all ranks skip). if is_ddp_initialized(): for v in scaler._found_inf_per_device(optimizer).values(): dist.all_reduce(v, op=dist.ReduceOp.MAX) scaler.step(optimizer) scaler.update() else: optimizer.step() model.zero_grad(set_to_none=True) train_loss_f = train_loss.item() # .item() is a CPU-GPU sync point synchronize() t1 = time.time() dt = t1 - t0 # ------------------------------------------------------------------------- # logging (CPU action only) ema_beta = 0.9 # EMA decay factor for some smoothing just for nicer logging smooth_train_loss = ema_beta * smooth_train_loss + (1 - ema_beta) * train_loss_f # 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 mfu = 100 * flops_per_sec / (gpu_peak_flops * ddp_world_size) if step > 10: total_training_time += dt # only count the time after the first 10 steps # Calculate ETA based on average time per step (excluding first 10 steps) steps_done = step - 10 if steps_done > 0: avg_time_per_step = total_training_time / steps_done remaining_steps = num_iterations - step eta_seconds = remaining_steps * avg_time_per_step eta_str = f" | eta: {eta_seconds/60:.1f}m" else: eta_str = "" epoch = f"{dataloader_state_dict['epoch']} pq: {dataloader_state_dict['pq_idx']} rg: {dataloader_state_dict['rg_idx']}" print0(f"step {step:05d}/{num_iterations:05d} ({pct_done:.2f}%) | loss: {debiased_smooth_loss:.6f} | lrm: {lrm:.2f} | dt: {dt * 1000:.2f}ms | tok/sec: {tok_per_sec:,} | bf16_mfu: {mfu:.2f} | epoch: {epoch} | total time: {total_training_time/60:.2f}m{eta_str}") 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, "train/epoch": epoch, } wandb_run.log(log_data) # state update first_step_of_run = (step == 0) or (resuming and step == args.resume_from_step) step += 1 # The garbage collector is sadly a little bit overactive and for some poorly understood reason, # it spends ~500ms scanning for cycles quite frequently, just to end up cleaning up very few tiny objects each time. # So we manually manage and help it out here if first_step_of_run: gc.collect() # manually collect a lot of garbage from setup gc.freeze() # immediately freeze all currently surviving objects and exclude them from GC gc.disable() # nuclear intervention here: disable GC entirely except: elif step % 5000 == 0: # every 5000 steps... gc.collect() # manually collect, just to be safe for very, very long runs # 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") if val_bpb is not None: print0(f"Minimum validation bpb: {min_val_bpb:.6f}") # 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 : Scaling params ratio": total_batch_size * num_iterations / num_scaling_params, "DDP world size": ddp_world_size, "warmup_steps": args.warmup_steps, "warmdown_ratio": args.warmdown_ratio, "final_lr_frac": args.final_lr_frac, }, { # stats about training outcomes "Minimum validation bpb": min_val_bpb if val_bpb is not None else None, "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() ================================================ FILE: scripts/chat_cli.py ================================================ """ 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 """ import argparse import torch from nanochat.common import compute_init, autodetect_device_type 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|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') 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) 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) 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 ================================================ FILE: scripts/chat_eval.py ================================================ """ Evaluate the Chat model. All the generic code lives here, and all the evaluation-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 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|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('-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) 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: 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() ================================================ FILE: scripts/chat_rl.py ================================================ """ 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 DAPO 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 argparse import os import itertools import wandb import torch import torch.distributed as dist from nanochat.common import compute_init, compute_cleanup, print0, get_base_dir, DummyWandb, autodetect_device_type from nanochat.checkpoint_manager import save_checkpoint, load_model from nanochat.engine import Engine from tasks.gsm8k import GSM8K # ----------------------------------------------------------------------------- # CLI arguments parser = argparse.ArgumentParser(description="Reinforcement learning on GSM8K") # Logging parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('dummy' disables wandb logging)") # Runtime parser.add_argument("--device-type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") # Model loading parser.add_argument("--model-tag", type=str, default=None, help="model tag to load from") parser.add_argument("--model-step", type=int, default=None, help="model step to load from") # Training horizon parser.add_argument("--num-epochs", type=int, default=1, help="number of epochs over GSM8K") # Batch sizes / sampling parser.add_argument("--device-batch-size", type=int, default=8, help="max batch size per forward pass") parser.add_argument("--examples-per-step", type=int, default=16, help="total examples per optimization step across all ranks") parser.add_argument("--num-samples", type=int, default=16, help="number of samples per example/question") # Generation parser.add_argument("--max-new-tokens", type=int, default=256, help="max tokens to generate per sample") parser.add_argument("--temperature", type=float, default=1.0, help="sampling temperature") parser.add_argument("--top-k", type=int, default=50, help="top-k sampling (0 = disabled)") # Optimization parser.add_argument("--embedding-lr", type=float, default=0.2, help="learning rate for embedding parameters (Adam)") parser.add_argument("--unembedding-lr", type=float, default=0.004, help="learning rate for unembedding parameters (Adam)") parser.add_argument("--matrix-lr", type=float, default=0.02, help="learning rate for matrix parameters (Muon)") parser.add_argument("--weight-decay", type=float, default=0.0, help="weight decay for embedding/unembedding parameters (Adam)") parser.add_argument("--init-lr-frac", type=float, default=0.05, help="initial LR as fraction of base LR") # Evaluation / checkpointing parser.add_argument("--eval-every", type=int, default=60, help="evaluate pass@k every N steps") parser.add_argument("--eval-examples", type=int, default=400, help="number of examples for pass@k evaluation") parser.add_argument("--save-every", type=int, default=60, help="save checkpoint every N steps") args = parser.parse_args() user_config = vars(args).copy() # ----------------------------------------------------------------------------- # Init compute/precision 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) master_process = ddp_rank == 0 # this process will do logging, checkpointing etc. # wandb logging init use_dummy_wandb = args.run == "dummy" or not master_process wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-rl", name=args.run, config=user_config) # Init model and tokenizer model, tokenizer, meta = load_model("sft", device, phase="eval", model_tag=args.model_tag, step=args.model_step) 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) // args.examples_per_step) * args.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 = args.num_samples // args.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 generated_token_sequences_batch, masks_batch = engine.generate_batch( tokens, num_samples=args.device_batch_size, max_tokens=args.max_new_tokens, temperature=args.temperature, top_k=args.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 <= args.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 optimizer = model.setup_optimizer( unembedding_lr=args.unembedding_lr, embedding_lr=args.embedding_lr, matrix_lr=args.matrix_lr, weight_decay=args.weight_decay, ) # Set the initial learning rate as a fraction of the base learning rate for group in optimizer.param_groups: group["lr"] = group["lr"] * args.init_lr_frac group["initial_lr"] = group["lr"] # 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: {args.examples_per_step * args.num_samples}") # total batch size in sequences/step assert args.examples_per_step % ddp_world_size == 0, "Desired examples per step must be divisible by the number of ranks" examples_per_rank = args.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 % args.eval_every == 0: model.eval() passk = torch.zeros(args.device_batch_size, device=device) # pass@k for k=1..device_batch_size records_iter = run_gsm8k_eval(val_task, tokenizer, engine, num_samples=args.device_batch_size, max_examples=args.eval_examples, temperature=1.0) records = list(records_iter) # collect all records for k in range(1, args.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, args.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, args.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) % args.device_batch_size == 0 num_passes = inputs_all.size(0) // args.device_batch_size for pass_idx in range(num_passes): # Pluck out the batch for this pass b0, b1 = pass_idx * args.device_batch_size, (pass_idx + 1) * args.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 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 group in optimizer.param_groups: group["lr"] = group["initial_lr"] * lrm optimizer.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 % args.save_every == 0) or step == num_steps - 1): base_dir = get_base_dir() depth = model.config.n_layer output_dirname = args.model_tag if args.model_tag else f"d{depth}" # base the model tag on the depth of the base model checkpoint_dir = os.path.join(base_dir, "chatrl_checkpoints", output_dirname) 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() ================================================ FILE: scripts/chat_sft.py ================================================ """ Supervised fine-tuning (SFT) the model. Run as: python -m scripts.chat_sft Or torchrun for training: torchrun --standalone --nproc_per_node=8 -m scripts.chat_sft -- --device-batch-size=16 """ import gc import argparse import os os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True" import time import wandb import torch from nanochat.common import compute_init, compute_cleanup, print0, DummyWandb, get_base_dir, autodetect_device_type, get_peak_flops, COMPUTE_DTYPE, COMPUTE_DTYPE_REASON, is_ddp_initialized from nanochat.tokenizer import get_token_bytes from nanochat.checkpoint_manager import save_checkpoint, load_model, load_optimizer_state from nanochat.loss_eval import evaluate_bpb import torch.distributed as dist from nanochat.flash_attention import HAS_FA3 from nanochat.engine import Engine from scripts.chat_eval import run_chat_eval 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 # ----------------------------------------------------------------------------- # CLI arguments parser = argparse.ArgumentParser(description="Supervised fine-tuning (SFT) the model") # Logging parser.add_argument("--run", type=str, default="dummy", help="wandb run name ('dummy' disables wandb logging)") # Runtime parser.add_argument("--device-type", type=str, default="", help="cuda|cpu|mps (empty = autodetect)") # Model loading parser.add_argument("--model-tag", type=str, default=None, help="model tag to load from") parser.add_argument("--model-step", type=int, default=None, help="model step to load from") parser.add_argument("--load-optimizer", type=int, default=1, help="warm-start optimizer from pretrained checkpoint (0=no, 1=yes)") # Training horizon parser.add_argument("--num-iterations", type=int, default=-1, help="number of optimization steps (-1 = full epoch)") # Batch sizes (default: inherit from pretrained checkpoint) parser.add_argument("--max-seq-len", type=int, default=None, help="max context length (default: inherit from pretrain)") parser.add_argument("--device-batch-size", type=int, default=None, help="per-device batch size (default: inherit from pretrain)") parser.add_argument("--total-batch-size", type=int, default=None, help="total batch size in tokens (default: inherit from pretrain)") # Optimization (default: inherit from pretrained checkpoint) parser.add_argument("--embedding-lr", type=float, default=None, help="learning rate for embedding parameters (Adam) (default: inherit from pretrain)") parser.add_argument("--unembedding-lr", type=float, default=None, help="learning rate for unembedding parameters (Adam) (default: inherit from pretrain)") parser.add_argument("--matrix-lr", type=float, default=None, help="learning rate for matrix parameters (Muon) (default: inherit from pretrain)") parser.add_argument("--init-lr-frac", type=float, default=0.8, help="initial LR as fraction of base LR") parser.add_argument("--warmup-ratio", type=float, default=0.0, help="ratio of iterations for LR warmup") parser.add_argument("--warmdown-ratio", type=float, default=0.5, help="ratio of iterations for LR warmdown") parser.add_argument("--final-lr-frac", type=float, default=0.0, help="final LR as fraction of initial LR") # Evaluation parser.add_argument("--eval-every", type=int, default=200, help="evaluate val bpb every N steps (-1 = disable)") parser.add_argument("--eval-tokens", type=int, default=40*524288, help="number of tokens to evaluate val loss on") parser.add_argument("--chatcore-every", type=int, default=200, help="evaluate ChatCORE metric every N steps (-1 = disable)") parser.add_argument("--chatcore-max-cat", type=int, default=-1, help="max problems per categorical task for ChatCORE") parser.add_argument("--chatcore-max-sample", type=int, default=24, help="max problems per generative task for ChatCORE") # Data mixture parser.add_argument("--mmlu-epochs", type=int, default=3, help="number of epochs of MMLU in training mixture (teaches Multiple Choice)") parser.add_argument("--gsm8k-epochs", type=int, default=4, help="number of epochs of GSM8K in training mixture (teaches Math and Tool Use)") args = parser.parse_args() user_config = vars(args).copy() # ----------------------------------------------------------------------------- # Compute init 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) master_process = ddp_rank == 0 print0(f"COMPUTE_DTYPE: {COMPUTE_DTYPE} ({COMPUTE_DTYPE_REASON})") 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 if device_type == "cuda": gpu_device_name = torch.cuda.get_device_name(0) gpu_peak_flops = get_peak_flops(gpu_device_name) print0(f"GPU: {gpu_device_name} | Peak FLOPS (BF16): {gpu_peak_flops:.2e}") else: gpu_peak_flops = float('inf') # MFU not meaningful for CPU/MPS # wandb logging init use_dummy_wandb = args.run == "dummy" or not master_process wandb_run = DummyWandb() if use_dummy_wandb else wandb.init(project="nanochat-sft", name=args.run, config=user_config) # Flash Attention status if not HAS_FA3: print0("WARNING: Flash Attention 3 not available, using PyTorch SDPA fallback. Training will be less efficient.") # Load the model and tokenizer model, tokenizer, meta = load_model("base", device, phase="train", model_tag=args.model_tag, step=args.model_step) # Inherit training hyperparameters from pretrained checkpoint (None = inherit, explicit value = override) pretrain_user_config = meta.get("user_config", {}) for name, fallback, source in [ ("max_seq_len", 2048, meta), ("device_batch_size", 32, meta), ("total_batch_size", 524288, meta), ("embedding_lr", 0.3, pretrain_user_config), ("unembedding_lr", 0.004, pretrain_user_config), ("matrix_lr", 0.02, pretrain_user_config), ]: arg_val = getattr(args, name) pretrain_val = source.get(name) if arg_val is None: resolved = pretrain_val if pretrain_val is not None else fallback setattr(args, name, resolved) print0(f"Inherited {name}={resolved} from pretrained checkpoint") elif pretrain_val is not None and arg_val != pretrain_val: print0(f"NOTE: --{name.replace('_', '-')}={arg_val} overrides pretrained value of {pretrain_val}") else: print0(f"Using {name}={arg_val}") orig_model = model model = torch.compile(model, dynamic=False) depth = model.config.n_layer num_flops_per_token = model.estimate_flops() tokens_per_fwdbwd = args.device_batch_size * args.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 args.total_batch_size % world_tokens_per_fwdbwd == 0 grad_accum_steps = args.total_batch_size // world_tokens_per_fwdbwd print0(f"Tokens / micro-batch / rank: {args.device_batch_size} x {args.max_seq_len} = {tokens_per_fwdbwd:,}") print0(f"Tokens / micro-batch: {world_tokens_per_fwdbwd:,}") print0(f"Total batch size {args.total_batch_size:,} => gradient accumulation steps: {grad_accum_steps}") token_bytes = get_token_bytes(device=device) # Initialize the Optimizer (combined MuonAdamW: Muon for matrix params, AdamW for rest) # Note that pretraining ramps weight_decay to zero by end of pretraining, so SFT continues with zero optimizer = model.setup_optimizer(unembedding_lr=args.unembedding_lr, embedding_lr=args.embedding_lr, matrix_lr=args.matrix_lr, weight_decay=0.0) # Optionally warm-start optimizer from pretrained checkpoint (momentum buffers etc.) # Note: load_state_dict overwrites param_group metadata (LRs, betas, etc.) with the # pretrained values. Since pretraining warmdown brings LRs to ~0, we must save and # restore our fresh SFT LRs after loading. base_dir = get_base_dir() if args.load_optimizer: optimizer_data = load_optimizer_state("base", device, rank=ddp_rank, model_tag=args.model_tag, step=args.model_step) if optimizer_data is not None: base_lrs = [group["lr"] for group in optimizer.param_groups] optimizer.load_state_dict(optimizer_data) del optimizer_data for group, base_lr in zip(optimizer.param_groups, base_lrs): group["lr"] = base_lr print0("Loaded optimizer state from pretrained checkpoint (momentum buffers only, LRs reset)") else: print0("WARNING: optimizer checkpoint not found, starting with fresh optimizer (slightly worse)") # GradScaler for fp16 training (bf16/fp32 don't need it) scaler = torch.amp.GradScaler() if COMPUTE_DTYPE == torch.float16 else None if scaler is not None: print0("GradScaler enabled for fp16 training") # Override the initial learning rate as a fraction of the base learning rate for group in optimizer.param_groups: group["lr"] = group["lr"] * args.init_lr_frac group["initial_lr"] = group["lr"] # SFT data mixture and DataLoader identity_conversations_filepath = os.path.join(base_dir, "identity_conversations.jsonl") train_tasks = [ SmolTalk(split="train"), # 460K rows of general conversations CustomJSON(filepath=identity_conversations_filepath), # 1000 rows of synthetic identity conversations CustomJSON(filepath=identity_conversations_filepath), # 2 epochs of these *[MMLU(subset="auxiliary_train", split="train") for _ in range(args.mmlu_epochs)], # 100K rows per epoch *[GSM8K(subset="main", split="train") for _ in range(args.gsm8k_epochs)], # 8K rows per epoch 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'?) ] train_dataset = TaskMixture(train_tasks) print0(f"Training mixture: {len(train_dataset):,} rows (MMLU x{args.mmlu_epochs}, GSM8K x{args.gsm8k_epochs})") 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 training dataset approx_progress = 0.0 # will go from 0 to 1 over the course of the epoch current_epoch = 1 # track epoch for logging def sft_data_generator_bos_bestfit(split, buffer_size=100): """ BOS-aligned dataloader for SFT with bestfit-pad packing. Each row in the batch starts with BOS (beginning of a conversation). Conversations are packed using best-fit algorithm. When no conversation fits, the row is padded (instead of cropping) to ensure no tokens are ever discarded. Padding positions have targets masked with -1 (ignore_index for cross-entropy). """ global last_step, approx_progress, current_epoch 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 row_capacity = args.max_seq_len + 1 # +1 for target at last position bos_token = tokenizer.get_bos_token_id() # Conversation buffer: list of (token_ids, loss_mask) tuples conv_buffer = [] cursor = ddp_rank # Each rank processes different conversations (for fetching) consumed = ddp_rank # Track actual consumption separately from buffering epoch = 1 it = 0 # iteration counter def refill_buffer(): nonlocal cursor, epoch while len(conv_buffer) < buffer_size: conversation = dataset[cursor] ids, mask = tokenizer.render_conversation(conversation) conv_buffer.append((ids, mask)) cursor += ddp_world_size if cursor >= dataset_size: cursor = cursor % dataset_size epoch += 1 # Note: last_step is now triggered based on consumption, not fetching while True: rows = [] mask_rows = [] row_lengths = [] # Track actual content length (excluding padding) for each row for _ in range(args.device_batch_size): row = [] mask_row = [] padded = False while len(row) < row_capacity: # Ensure buffer has conversations while len(conv_buffer) < buffer_size: refill_buffer() remaining = row_capacity - len(row) # Find largest conversation that fits entirely best_idx = -1 best_len = 0 for i, (conv, _) in enumerate(conv_buffer): conv_len = len(conv) if conv_len <= remaining and conv_len > best_len: best_idx = i best_len = conv_len if best_idx >= 0: # Found a conversation that fits - use it entirely conv, conv_mask = conv_buffer.pop(best_idx) row.extend(conv) mask_row.extend(conv_mask) consumed += ddp_world_size # Track actual consumption else: # No conversation fits - pad the remainder instead of cropping # This ensures we never discard any tokens content_len = len(row) row.extend([bos_token] * remaining) # Pad with BOS tokens mask_row.extend([0] * remaining) padded = True break # Row is now full (with padding) # Track content length: full row if no padding, otherwise the length before padding if padded: row_lengths.append(content_len) else: row_lengths.append(row_capacity) rows.append(row[:row_capacity]) mask_rows.append(mask_row[:row_capacity]) # Stopping condition to respect num_iterations, if given it += 1 if 0 < args.num_iterations <= it and split == "train": last_step = True # Update progress tracking (based on consumed, not cursor, to account for buffering) if split == "train": current_epoch = epoch if args.num_iterations > 0: approx_progress = it / args.num_iterations else: approx_progress = consumed / dataset_size # Trigger last_step when we've consumed enough (instead of when cursor wraps) if consumed >= dataset_size: last_step = True # Build tensors use_cuda = device_type == "cuda" batch_tensor = torch.tensor(rows, dtype=torch.long, pin_memory=use_cuda) inputs = batch_tensor[:, :-1].to(device=device, dtype=torch.int32, non_blocking=use_cuda).contiguous() targets = batch_tensor[:, 1:].to(device=device, dtype=torch.int64, non_blocking=use_cuda).contiguous() # Apply the loss mask from render_conversation (mask=1 for assistant completions, # mask=0 for user prompts, BOS, special tokens, tool outputs). mask[1:] aligns # with targets (shifted by 1). Unmasked positions get -1 (ignore_index). mask_tensor = torch.tensor(mask_rows, dtype=torch.int8) mask_targets = mask_tensor[:, 1:].to(device=device) targets[mask_targets == 0] = -1 # Mask out padding positions in targets (set to -1 = ignore_index) # For each row, positions >= (content_length - 1) in targets should be masked for i, content_len in enumerate(row_lengths): if content_len < row_capacity: targets[i, content_len-1:] = -1 yield inputs, targets train_loader = sft_data_generator_bos_bestfit("train") build_val_loader = lambda: sft_data_generator_bos_bestfit("val") progress = 0 # will go from 0 to 1 over the course of the epoch # Learning rate schedule (linear warmup, constant, linear warmdown) # Same shape as base_train but uses progress (0→1) instead of absolute step counts, # because SFT doesn't always know num_iterations in advance (dataset-driven stopping). def get_lr_multiplier(progress): if progress < args.warmup_ratio: return (progress + 1e-8) / args.warmup_ratio elif progress <= 1.0 - args.warmdown_ratio: return 1.0 else: decay = (progress - (1.0 - args.warmdown_ratio)) / args.warmdown_ratio return (1 - decay) * 1.0 + decay * args.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 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 * args.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 last_step or (args.eval_every > 0 and step % args.eval_every == 0): model.eval() val_loader = build_val_loader() eval_steps = args.eval_tokens // (args.device_batch_size * args.max_seq_len * ddp_world_size) 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 ChatCORE metric (all ranks participate) # use the original uncompiled model because the inputs keep changing shape chatcore_results = {} if args.chatcore_every > 0 and (last_step or (step > 0 and step % args.chatcore_every == 0)): model.eval() engine = Engine(orig_model, tokenizer) all_tasks = ['ARC-Easy', 'ARC-Challenge', 'MMLU', 'GSM8K', 'HumanEval', 'SpellingBee'] categorical_tasks = {'ARC-Easy', 'ARC-Challenge', 'MMLU'} baseline_accuracies = { 'ARC-Easy': 0.25, 'ARC-Challenge': 0.25, 'MMLU': 0.25, 'GSM8K': 0.0, 'HumanEval': 0.0, 'SpellingBee': 0.0, } task_results = {} for task_name in all_tasks: limit = args.chatcore_max_cat if task_name in categorical_tasks else args.chatcore_max_sample max_problems = None if limit < 0 else limit # -1 means no limit acc = run_chat_eval(task_name, orig_model, tokenizer, engine, batch_size=args.device_batch_size, max_problems=max_problems) task_results[task_name] = acc print0(f" {task_name}: {100*acc:.2f}%") # Compute ChatCORE metrics (mean centered accuracy, ranges from 0=random to 1=perfect) def centered_mean(tasks): return sum((task_results[t] - baseline_accuracies[t]) / (1.0 - baseline_accuracies[t]) for t in tasks) / len(tasks) chatcore = centered_mean(all_tasks) chatcore_cat = centered_mean(categorical_tasks) print0(f"Step {step:05d} | ChatCORE: {chatcore:.4f} | ChatCORE_cat: {chatcore_cat:.4f}") wandb_run.log({ "step": step, "total_training_flops": flops_so_far, "chatcore_metric": chatcore, "chatcore_cat": chatcore_cat, **{f"chatcore/{task_name}": acc for task_name, acc in task_results.items()}, }) model.train() # save checkpoint at the end of the run (all ranks participate so each saves its optimizer shard) if last_step: output_dirname = args.model_tag if args.model_tag else f"d{depth}" # e.g. d12 checkpoint_dir = os.path.join(base_dir, "chatsft_checkpoints", output_dirname) save_checkpoint( checkpoint_dir, step, orig_model.state_dict(), optimizer.state_dict(), { "step": step, "val_bpb": val_bpb, # loss at last step "model_config": { "sequence_len": args.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, "window_pattern": model.config.window_pattern, }, "user_config": user_config, # inputs to the training script }, rank=ddp_rank, ) if last_step: break # ------------------------------------------------------------------------- # single training step # evaluate the gradient synchronize() t0 = time.time() for micro_step in range(grad_accum_steps): loss = model(x, y) train_loss = loss.detach() # for logging loss = loss / grad_accum_steps # each .backward() is a grad sum => normalize loss here if scaler is not None: scaler.scale(loss).backward() else: 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 optimizer lrm = get_lr_multiplier(progress) muon_momentum = get_muon_momentum(step) for group in optimizer.param_groups: group["lr"] = group["initial_lr"] * lrm if group['kind'] == 'muon': group["momentum"] = muon_momentum if scaler is not None: scaler.unscale_(optimizer) if is_ddp_initialized(): for v in scaler._found_inf_per_device(optimizer).values(): dist.all_reduce(v, op=dist.ReduceOp.MAX) scaler.step(optimizer) scaler.update() else: optimizer.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(args.total_batch_size / dt) flops_per_sec = num_flops_per_token * args.total_batch_size / dt mfu = 100 * flops_per_sec / (gpu_peak_flops * ddp_world_size) 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} | epoch: {current_epoch} | 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, "train/epoch": current_epoch, }) # The garbage collector spends ~500ms scanning for cycles quite frequently. # We manually manage it to avoid these pauses during training. if step == 1: gc.collect() # manually collect a lot of garbage from setup gc.freeze() # freeze all currently surviving objects and exclude them from GC gc.disable() # disable GC entirely except: elif step % 5000 == 0: # every 5000 steps... gc.collect() # manually collect, just to be safe for very long runs # 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="SFT", 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() ================================================ FILE: scripts/chat_web.py ================================================ #!/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 0-200 (0 disables top-k filtering, using full vocabulary) - 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 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 = 0 # 0 disables top-k filtering, using full vocabulary 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|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('--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) @dataclass class Worker: """A worker with a model loaded on a specific GPU.""" gpu_id: int device: torch.device engine: Engine tokenizer: object 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) worker = Worker( gpu_id=gpu_id, device=device, engine=engine, tokenizer=tokenizer, ) 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 = "" 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) ================================================ FILE: scripts/tok_eval.py ================================================ """ 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, ]) ================================================ FILE: scripts/tok_train.py ================================================ """ Train a tokenizer using our own BPE Tokenizer 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=2_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=32768, help='Vocabulary size (default: 32768 = 2^15)') 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(), } ]) ================================================ FILE: tasks/arc.py ================================================ """ 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 ================================================ FILE: tasks/common.py ================================================ """ 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 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]) ================================================ FILE: tasks/customjson.py ================================================ """ 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 suddenly 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 ================================================ FILE: tasks/gsm8k.py ================================================ """ 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}) # Now 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 ================================================ FILE: tasks/humaneval.py ================================================ """ Evaluate the Chat model on HumanEval dataset. Btw this dataset is a misnomer and has nothing to do with humans. It is a coding benchmark. """ import re from datasets import load_dataset from nanochat.execution import execute_code from tasks.common import Task def extract_imports(prompt): """Extract import statements from the beginning of a code block.""" imports = [] for line in prompt.split('\n'): stripped = line.strip() if stripped.startswith('import ') or stripped.startswith('from '): imports.append(stripped) elif stripped and not stripped.startswith('#'): # Stop at first non-import, non-comment line break return '\n'.join(imports) def extract_program(completion): """ Extract Python code from LLM completion. Handles various output formats: - Code wrapped in ```python ... ``` or ``` ... ``` blocks - Plain code without markdown blocks - Extra text before/after code blocks Returns the first code block if found, otherwise returns the whole completion. """ # Try to find markdown code blocks (```python or just ```) # Match ```python\n...\n``` or ```\n...\n``` pattern = r'```(?:python)?\s*\n(.*?)\n```' matches = re.findall(pattern, completion, re.DOTALL) if matches: # Return the first code block found return matches[0].strip() # No code blocks found, return the whole completion return completion.strip() class HumanEval(Task): def __init__(self, **kwargs): super().__init__(**kwargs) self.ds = load_dataset("openai/openai_humaneval", split="test").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] prompt = row['prompt'] # prompts in HumanEval are the beginning of the program solution = row['canonical_solution'] # the correct continuation of the program entry_point = row['entry_point'] # the function to check test = row['test'] # the test cases complete_solution = f"{prompt}\n{solution}" messages = [ {"role": "user", "content": prompt}, {"role": "assistant", "content": complete_solution}, ] conversation = { "messages": messages, "entry_point": entry_point, # needed during evaluation "test": test, # needed during evaluation } return conversation def evaluate(self, conversation, completion): """ Given (conversation, completion), return boolean success of the completion. """ # the prompt will contain the imports and the function signature imports = extract_imports(conversation['messages'][0]['content']) # the completion will usually contain the whole function # but not always with the needed imports, so we manually append them completion_code = extract_program(completion) program = ( imports + "\n\n" + completion_code + "\n\n" + conversation['test'] + "\n" + f"check({conversation['entry_point']})" ) result = execute_code(program) success = result.success return success ================================================ FILE: tasks/mmlu.py ================================================ """ The MMLU dataset. https://huggingface.co/datasets/cais/mmlu """ from datasets import load_dataset from tasks.common import Task, render_mc class MMLU(Task): letters = ('A', 'B', 'C', 'D') groups = ('abstract_algebra', 'anatomy', 'astronomy', 'business_ethics', 'clinical_knowledge', 'college_biology', 'college_chemistry', 'college_computer_science', 'college_mathematics', 'college_medicine', 'college_physics', 'computer_security', 'conceptual_physics', 'econometrics', 'electrical_engineering', 'elementary_mathematics', 'formal_logic', 'global_facts', 'high_school_biology', 'high_school_chemistry', 'high_school_computer_science', 'high_school_european_history', 'high_school_geography', 'high_school_government_and_politics', 'high_school_macroeconomics', 'high_school_mathematics', 'high_school_microeconomics', 'high_school_physics', 'high_school_psychology', 'high_school_statistics', 'high_school_us_history', 'high_school_world_history', 'human_aging', 'human_sexuality', 'international_law', 'jurisprudence', 'logical_fallacies', 'machine_learning', 'management', 'marketing', 'medical_genetics', 'miscellaneous', 'moral_disputes', 'moral_scenarios', 'nutrition', 'philosophy', 'prehistory', 'professional_accounting', 'professional_law', 'professional_medicine', 'professional_psychology', 'public_relations', 'security_studies', 'sociology', 'us_foreign_policy', 'virology', 'world_religions') def __init__(self, subset, split, **kwargs): super().__init__(**kwargs) assert subset in ["all", "auxiliary_train"], f"subset {subset} must be all|auxiliary_train" assert split in ["train", "validation", "dev", "test"], f"split {split} must be train|validation|dev|test" if subset == "auxiliary_train": assert split == "train", "auxiliary_train must be split into train" self.subset = subset self.split = split self.ds = load_dataset("cais/mmlu", subset, split=split).shuffle(seed=42) if subset == "auxiliary_train": # I don't understand why but the auxiliary_train rows have some weird additional 'train' wrapper self.ds = self.ds.map(lambda row: row['train'], remove_columns=['train']) @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"] # the text of each choice answer = row["answer"] # index of the answer, e.g. 0,1,2,3 (for A,B,C,D) subject = row["subject"] # e.g. "college_biology", "college_chemistry", etc. assert len(choices) == 4, "MMLU should have 4 choices" # create and return the Conversation object user_message = render_mc(question, self.letters, choices) assistant_message = self.letters[answer] messages = [ {"role": "user", "content": user_message}, {"role": "assistant", "content": assistant_message} ] conversation = { "messages": messages, "subject": subject, # might be useful later for grouping metrics by subject "letters": self.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 self.letters, f"MMLU answer {assistant_response} is expected to be one of {self.letters}" assistant_message = conversation['messages'][-1]['content'] # e.g. "A" return assistant_response == assistant_message ================================================ FILE: tasks/smoltalk.py ================================================ """ SmolTalk by HuggingFace. Good "general" conversational dataset. https://huggingface.co/datasets/HuggingFaceTB/smol-smoltalk We use the "smol" version, which is more appropriate for smaller models. """ from datasets import load_dataset from tasks.common import Task class SmolTalk(Task): """ smol-smoltalk dataset. train is 460K rows, test is 24K rows. """ def __init__(self, split, **kwargs): super().__init__(**kwargs) assert split in ["train", "test"], "SmolTalk split must be train|test" self.ds = load_dataset("HuggingFaceTB/smol-smoltalk", split=split).shuffle(seed=42) self.length = len(self.ds) def num_examples(self): return self.length def get_example(self, index): row = self.ds[index] messages = row["messages"] # --------------------------------------------------------------------- # sanity checking asserts here # TODO: we could remove these asserts later, for now just don't want any footguns # there is an optional system message at the beginning assert len(messages) >= 1 first_message = messages[0] if first_message["role"] == "system": rest_messages = messages[1:] # optional system message is OK else: rest_messages = messages assert len(rest_messages) >= 2, "SmolTalk messages must have at least 2 messages" for i, message in enumerate(rest_messages): # user and assistant alternate as user,assistant,user,assistant,... 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), "Content must be a string" # --------------------------------------------------------------------- # create and return the Conversation object (ok to emit the system message too) conversation = { "messages": messages, } return conversation ================================================ FILE: tasks/spellingbee.py ================================================ """ Task intended to make nanochat better in spelling and counting, for example: "How many r are in strawberry?" -> 3 An interesting part of this task is that we will get the assistant to solve the problem using a combination of manual counting and Python. This is a good problem solving "instinct" to mix into the model and RL may further refine it to trust one over the other. If we were extra fancy (which we could/should be) we'd add small errors here and there to allow the model also learn recoveries. We can do this in future versions. There are two tasks in this file: 1. SpellingBee: Counting the number of occurrences of a letter in a word 2. SimpleSpelling: Simply spelling words (1) is the goal, but (2) exists as a highly condensed version of the part that makes (1) difficult, which is word spelling. This is non-trivial for an LLM because it has to learn how every token (a little semantic chunk/atom) maps to the sequence of individual characters that make it up. Larger models learn this eventually on their own, but if we want this capability to exist in smaller models, we have to actively encourage it by over-representing it in the training data. SFT is a good place to do this. To preview a few example conversations, run: python -m tasks.spellingbee """ import re import random from tasks.common import Task from nanochat.common import download_file_with_lock # Letters of the alphabet LETTERS = "abcdefghijklmnopqrstuvwxyz" # A list of 370K English words of large variety WORD_LIST_URL = "https://raw.githubusercontent.com/dwyl/english-words/refs/heads/master/words_alpha.txt" # A number bigger than 370K to separate train and test random seeds TEST_RANDOM_SEED_OFFSET = 10_000_000 # Identical to gsm8k's answer extraction ANSWER_RE = re.compile(r"#### (\-?[0-9\.\,]+)") def extract_answer(completion): """ Extract the numerical answer after #### marker. """ match = ANSWER_RE.search(completion) if match: match_str = match.group(1).strip() match_str = match_str.replace(",", "") return match_str return None # User message templates for data augmentation USER_MSG_TEMPLATES = [ "How many {letter} are in the word {word}", "How many {letter} are in {word}", "Count the number of {letter} in {word}", "How many times does {letter} appear in {word}", "What's the count of {letter} in {word}", "In the word {word}, how many {letter} are there", "How many letter {letter} are in the word {word}", "Count how many {letter} appear in {word}", "Tell me the number of {letter} in {word}", "How many occurrences of {letter} are in {word}", "Find the count of {letter} in {word}", "Can you count the {letter} letters in {word}", "What is the frequency of {letter} in {word}", "How many {letter}s are in {word}", "How many {letter}'s are in {word}", "Count all the {letter} in {word}", "How many times is {letter} in {word}", "Number of {letter} in {word}", "Total count of {letter} in {word}", "How many {letter} does {word} have", "How many {letter} does {word} contain", "What's the number of {letter} in {word}", "{word} has how many {letter}", "In {word}, count the {letter}", "How many {letter} appear in {word}", "Count the {letter} in {word}", "Give me the count of {letter} in {word}", "How many instances of {letter} in {word}", "Show me how many {letter} are in {word}", "Calculate the number of {letter} in {word}", # Spanish "¿Cuántas {letter} hay en {word}?", "¿Cuántas veces aparece {letter} en {word}?", "Cuenta las {letter} en {word}", "¿Cuántas letras {letter} tiene {word}?", # Chinese (Simplified) "{word}中有多少个{letter}", "{word}里有几个{letter}", "数一下{word}中的{letter}", "{word}这个词里有多少{letter}", # Korean "{word}에 {letter}가 몇 개 있나요", "{word}에서 {letter}의 개수는", "{word}에 {letter}가 몇 번 나오나요", "{word}라는 단어에 {letter}가 몇 개", # French "Combien de {letter} dans {word}", "Combien de fois {letter} apparaît dans {word}", "Compte les {letter} dans {word}", # German "Wie viele {letter} sind in {word}", "Wie oft kommt {letter} in {word} vor", "Zähle die {letter} in {word}", # Japanese "{word}に{letter}は何個ありますか", "{word}の中に{letter}がいくつ", "{word}に{letter}が何回出てくる", ] class SpellingBee(Task): def __init__(self, size=1000, split="train", **kwargs): super().__init__(**kwargs) assert split in ["train", "test"], "SpellingBee split must be train|test" self.size = size self.split = split filename = WORD_LIST_URL.split("/")[-1] word_list_path = download_file_with_lock(WORD_LIST_URL, filename) with open(word_list_path, 'r', encoding='utf-8') as f: words = [line.strip() for line in f] self.words = words @property def eval_type(self): return 'generative' def num_examples(self): return self.size def get_example(self, index): seed = index if self.split == 'train' else TEST_RANDOM_SEED_OFFSET + index rng = random.Random(seed) # pick a random word word = rng.choice(self.words) # pick a letter from it (90%) or a random letter (10%) letter = rng.choice(word) if rng.random() < 0.9 else rng.choice(LETTERS) # get the correct answer by simply counting count = word.count(letter) # create a user message, with a bunch of variations as data augmentation template = rng.choice(USER_MSG_TEMPLATES) # 30% chance to lowercase the template (lazy people don't use shift) if rng.random() < 0.3: template = template.lower() quote_options = ['', "'", '"'] letter_quote = rng.choice(quote_options) # is the letter quoted? word_quote = rng.choice(quote_options) # is the word quoted? letter_wrapped = f"{letter_quote}{letter}{letter_quote}" word_wrapped = f"{word_quote}{word}{word_quote}" user_msg = template.format(letter=letter_wrapped, word=word_wrapped) if rng.random() < 0.5: # 50% of people don't even use question marks user_msg += "?" # Now create the ideal assistant response - build as parts (text + tool calls) assistant_parts = [] word_letters = ",".join(list(word)) manual_text = f"""We are asked to find the number '{letter}' in the word '{word}'. Let me try a manual approach first. First spell the word out: {word}:{word_letters} Then count the occurrences of '{letter}': """ # Little simulated loop of the solution process # TODO: This is where the fun starts, we could simulate cute little mistakes # and get the model to review its work and recover from them. # You might of course hope this could arise in RL too, but realistically you'd want to help it out a bit. running_count = 0 for i, char in enumerate(word, 1): if char == letter: running_count += 1 # note: there deliberately cannot be a space here between i and char # because this would create a different token! (e.g. " a" and "a" are different tokens) manual_text += f"{i}:{char} hit! count={running_count}\n" else: manual_text += f"{i}:{char}\n" manual_text += f"\nThis gives us {running_count}." assistant_parts.append({"type": "text", "text": manual_text}) # Part 2: Python verification assistant_parts.append({"type": "text", "text": "\n\nLet me double check this using Python:\n\n"}) # Part 3: Python tool call python_expr = f"'{word}'.count('{letter}')" assistant_parts.append({"type": "python", "text": python_expr}) # Part 4: Python output assistant_parts.append({"type": "python_output", "text": str(count)}) # Part 5: Final answer assistant_parts.append({"type": "text", "text": f"\n\nPython gives us {count}.\n\nMy final answer is:\n\n#### {count}"}) # return the full conversation messages = [ {"role": "user", "content": user_msg}, {"role": "assistant", "content": assistant_parts} ] conversation = { "messages": messages, } return conversation def evaluate(self, conversation, assistant_response): """ Given (conversation, completion), return evaluation outcome (0 = wrong, 1 = correct) Identical to gsm8k's evaluation. """ assert isinstance(assistant_response, str), "Assuming simple string response for now" # First extract the ground truth answer from the conversation 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" # The last text part contains the final answer with #### last_text_part = assistant_message['content'][-1]['text'] # 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): """ Use simple 0-1 reward just like gsm8k.""" is_correct = self.evaluate(conversation, assistant_response) is_correct_float = float(is_correct) return is_correct_float class SimpleSpelling(Task): """Much simpler task designed to get the model to just practice spelling words.""" def __init__(self, size=1000, split="train", **kwargs): super().__init__(**kwargs) assert split in ["train", "test"], "SpellingBee split must be train|test" self.size = size self.split = split filename = WORD_LIST_URL.split("/")[-1] word_list_path = download_file_with_lock(WORD_LIST_URL, filename) with open(word_list_path, 'r', encoding='utf-8') as f: words = [line.strip() for line in f] rng = random.Random(42) rng.shuffle(words) # use a different word order than the SpellingBee task self.words = words @property def eval_type(self): return 'generative' def num_examples(self): return self.size def get_example(self, index): seed = index if self.split == 'train' else TEST_RANDOM_SEED_OFFSET + index rng = random.Random(seed) # pick a random word word = rng.choice(self.words) word_letters = ",".join(list(word)) # return the full conversation messages = [ {"role": "user", "content": f"Spell the word: {word}"}, {"role": "assistant", "content": f"{word}:{word_letters}"} ] conversation = { "messages": messages, } return conversation if __name__ == "__main__": # preview the SpellingBee task, first 10 examples task = SpellingBee() for i in range(10): ex = task.get_example(i) print("=" * 100) print(ex['messages'][0]['content']) print("-" * 100) # Assistant content is now a list of parts assistant_parts = ex['messages'][1]['content'] for part in assistant_parts: if part['type'] == 'text': print(part['text'], end='') elif part['type'] == 'python': print(f"<<{part['text']}=", end='') elif part['type'] == 'python_output': print(f"{part['text']}>>", end='') print() print("-" * 100) # # preview the SimpleSpelling task, first 10 examples # task = SimpleSpelling() # for i in range(10): # ex = task.get_example(i) # print("=" * 100) # print(ex['messages'][0]['content']) # print("-" * 100) # print(ex['messages'][1]['content']) # # also scrutinize the tokenization (last example only) # from nanochat.tokenizer import get_tokenizer # tokenizer = get_tokenizer() # ids, mask = tokenizer.render_conversation(ex) # print(tokenizer.visualize_tokenization(ids, mask, with_token_id=True)) ================================================ FILE: tests/test_attention_fallback.py ================================================ """ Test Flash Attention unified interface - verify FA3 and SDPA produce identical results. Run: python -m pytest tests/test_attention_fallback.py -v -s Note on test structure: Tests are split into two classes due to dtype/device constraints: 1. TestFA3VsSDPA: Comparison tests that run both FA3 and SDPA on the same inputs and verify they produce identical results. These require a Hopper GPU (FA3 only works on sm90+) and use bfloat16 (FA3 doesn't support float32). 2. TestSDPAOnly: Tests that only exercise the SDPA fallback path. These can run on any device (CUDA, CPU, MPS) with the appropriate dtype for that device. """ import torch import pytest import nanochat.flash_attention as fa_module from nanochat.flash_attention import flash_attn, HAS_FA3 from nanochat.engine import KVCache def set_impl(impl): """Set the implementation override ('fa3', 'sdpa', or None for auto) and re-resolve USE_FA3.""" fa_module._override_impl = impl fa_module.USE_FA3 = fa_module._resolve_use_fa3() def run_both_impls(fn): """Run a function with both FA3 and SDPA, return both outputs.""" set_impl('fa3') out_fa3 = fn() set_impl('sdpa') out_sdpa = fn() set_impl(None) # reset return out_fa3, out_sdpa def assert_close(t1, t2, name, atol=1e-2, rtol=1e-2): """Assert two tensors are close, with helpful error message.""" max_diff = (t1 - t2).abs().max().item() mean_diff = (t1 - t2).abs().mean().item() assert torch.allclose(t1, t2, atol=atol, rtol=rtol), \ f"{name}: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}" return max_diff, mean_diff # ============================================================================= # FA3 vs SDPA comparison tests (require Hopper GPU) # ============================================================================= @pytest.mark.skipif(not HAS_FA3, reason="FA3 required to compare implementations") class TestFA3VsSDPA: """Compare FA3 and SDPA produce identical results. Requires Hopper GPU.""" DEVICE = "cuda" DTYPE = torch.bfloat16 def test_basic_causal(self): """Basic causal attention.""" B, T, H, D = 2, 64, 4, 32 q = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): return flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(T, 0)) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "basic_causal") print(f"basic_causal: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_full_context(self): """Full context (window_size=-1).""" B, T, H, D = 2, 128, 4, 32 q = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): return flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(-1, -1)) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "full_context") print(f"full_context: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_sliding_window(self): """Sliding window attention.""" B, T, H, D = 2, 128, 4, 32 window = 32 q = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): return flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(window, 0)) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "sliding_window") print(f"sliding_window: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_gqa(self): """Group Query Attention (fewer KV heads than Q heads).""" B, T, D = 2, 64, 32 n_heads = 8 n_kv_heads = 2 q = torch.randn(B, T, n_heads, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T, n_kv_heads, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T, n_kv_heads, D, device=self.DEVICE, dtype=self.DTYPE) def run(): return flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(T, 0)) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "gqa") print(f"gqa: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_larger_model(self): """Larger dimensions closer to real model.""" B, T, H, D = 4, 256, 12, 64 q = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): return flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(-1, -1)) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "larger_model") print(f"larger_model: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_kvcache_prefill(self): """Test prefill (inserting multiple tokens into empty cache).""" B, T_max, H, D = 2, 64, 4, 32 T_prefill = 16 q = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): k_cache = torch.zeros(B, T_max, H, D, device=self.DEVICE, dtype=self.DTYPE) v_cache = torch.zeros(B, T_max, H, D, device=self.DEVICE, dtype=self.DTYPE) cache_seqlens = torch.zeros(B, dtype=torch.int32, device=self.DEVICE) return flash_attn.flash_attn_with_kvcache( q, k_cache, v_cache, k=k, v=v, cache_seqlens=cache_seqlens, causal=True, window_size=(T_max, 0) ) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "prefill") print(f"prefill: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_kvcache_single_token(self): """Test single token generation (cache already has content).""" B, T_max, H, D = 2, 64, 4, 32 T_prefill = 16 k_init = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) v_init = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) q_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) k_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) v_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): k_cache = torch.zeros(B, T_max, H, D, device=self.DEVICE, dtype=self.DTYPE) v_cache = torch.zeros(B, T_max, H, D, device=self.DEVICE, dtype=self.DTYPE) k_cache[:, :T_prefill, :, :] = k_init v_cache[:, :T_prefill, :, :] = v_init cache_seqlens = torch.full((B,), T_prefill, dtype=torch.int32, device=self.DEVICE) return flash_attn.flash_attn_with_kvcache( q_single, k_cache, v_cache, k=k_single, v=v_single, cache_seqlens=cache_seqlens, causal=True, window_size=(T_max, 0) ) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "single_token") print(f"single_token: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_kvcache_single_token_sliding_window(self): """Test single token decode with sliding window smaller than cache size. This catches the bug where SDPA ignores window_size during Tq=1 decode. When window < Tk, FA3 only attends to the last (window+1) tokens, but SDPA was attending to all cached tokens. """ B, T_max, H, D = 2, 64, 4, 32 T_prefill = 32 # Enough tokens to exceed window window = 8 # Window SMALLER than cache size k_init = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) v_init = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) q_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) k_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) v_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): k_cache = torch.zeros(B, T_max, H, D, device=self.DEVICE, dtype=self.DTYPE) v_cache = torch.zeros(B, T_max, H, D, device=self.DEVICE, dtype=self.DTYPE) k_cache[:, :T_prefill, :, :] = k_init v_cache[:, :T_prefill, :, :] = v_init cache_seqlens = torch.full((B,), T_prefill, dtype=torch.int32, device=self.DEVICE) return flash_attn.flash_attn_with_kvcache( q_single, k_cache, v_cache, k=k_single, v=v_single, cache_seqlens=cache_seqlens, causal=True, window_size=(window, 0) # window=8 < Tk=33 ) y_fa3, y_sdpa = run_both_impls(run) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "single_token_sliding_window") print(f"single_token_sliding_window: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") def test_backward_gradients_match(self): """Verify gradients are similar between FA3 and SDPA.""" B, T, H, D = 2, 32, 4, 16 q_data = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) k_data = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) v_data = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) def run(): q = q_data.clone().requires_grad_(True) k = k_data.clone().requires_grad_(True) v = v_data.clone().requires_grad_(True) y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(T, 0)) loss = y.sum() loss.backward() return y.detach(), q.grad.detach(), k.grad.detach(), v.grad.detach() set_impl('fa3') y_fa3, q_grad_fa3, k_grad_fa3, v_grad_fa3 = run() set_impl('sdpa') y_sdpa, q_grad_sdpa, k_grad_sdpa, v_grad_sdpa = run() set_impl(None) max_diff, mean_diff = assert_close(y_fa3, y_sdpa, "backward_output") print(f"backward_output: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") max_diff, mean_diff = assert_close(q_grad_fa3, q_grad_sdpa, "q_grad", atol=0.05, rtol=0.05) print(f"q_grad: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") max_diff, mean_diff = assert_close(k_grad_fa3, k_grad_sdpa, "k_grad", atol=0.05, rtol=0.05) print(f"k_grad: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") max_diff, mean_diff = assert_close(v_grad_fa3, v_grad_sdpa, "v_grad", atol=0.05, rtol=0.05) print(f"v_grad: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}") # ============================================================================= # SDPA-only tests (run on any device) # ============================================================================= class TestSDPAOnly: """Test SDPA fallback works correctly. Runs on any device.""" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32 def test_basic_forward(self): """Test SDPA forward pass produces valid output.""" set_impl('sdpa') B, T, H, D = 2, 64, 4, 32 q = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE) y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(T, 0)) assert y.shape == (B, T, H, D) assert not torch.isnan(y).any(), "Output contains NaN" set_impl(None) def test_backward(self): """Test gradients flow through SDPA.""" set_impl('sdpa') B, T, H, D = 2, 32, 4, 16 q = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE, requires_grad=True) k = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE, requires_grad=True) v = torch.randn(B, T, H, D, device=self.DEVICE, dtype=self.DTYPE, requires_grad=True) y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=(T, 0)) loss = y.sum() loss.backward() assert q.grad is not None, "No gradient for q" assert k.grad is not None, "No gradient for k" assert v.grad is not None, "No gradient for v" assert not torch.isnan(q.grad).any(), "NaN in q gradient" set_impl(None) def test_kvcache(self): """Test SDPA with KV cache.""" set_impl('sdpa') B, T_max, H, D = 2, 64, 4, 32 n_layers = 1 cache = KVCache( batch_size=B, num_heads=H, seq_len=T_max, head_dim=D, num_layers=n_layers, device=self.DEVICE, dtype=self.DTYPE ) k_cache, v_cache = cache.get_layer_cache(0) # Prefill T_prefill = 16 q = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) k = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) v = torch.randn(B, T_prefill, H, D, device=self.DEVICE, dtype=self.DTYPE) y = flash_attn.flash_attn_with_kvcache( q, k_cache, v_cache, k=k, v=v, cache_seqlens=cache.cache_seqlens, causal=True, window_size=(T_max, 0) ) cache.advance(T_prefill) assert y.shape == (B, T_prefill, H, D) assert cache.get_pos() == T_prefill # Generate single token q_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) k_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) v_single = torch.randn(B, 1, H, D, device=self.DEVICE, dtype=self.DTYPE) y_single = flash_attn.flash_attn_with_kvcache( q_single, k_cache, v_cache, k=k_single, v=v_single, cache_seqlens=cache.cache_seqlens, causal=True, window_size=(T_max, 0) ) cache.advance(1) assert y_single.shape == (B, 1, H, D) assert cache.get_pos() == T_prefill + 1 set_impl(None) # ============================================================================= # Override mechanism tests # ============================================================================= class TestOverrideMechanism: """Test that the override mechanism works correctly.""" @pytest.mark.skipif(not HAS_FA3, reason="FA3 required") def test_override_fa3(self): """Test that override='fa3' uses FA3.""" set_impl('fa3') assert fa_module.USE_FA3 == True set_impl(None) def test_override_sdpa(self): """Test that override='sdpa' uses SDPA.""" set_impl('sdpa') assert fa_module.USE_FA3 == False set_impl(None) def test_override_auto(self): """Test that override=None uses auto-detection.""" set_impl(None) assert fa_module.USE_FA3 == HAS_FA3 if __name__ == "__main__": print(f"PyTorch version: {torch.__version__}") print(f"CUDA available: {torch.cuda.is_available()}") if torch.cuda.is_available(): print(f"CUDA device: {torch.cuda.get_device_name()}") major, minor = torch.cuda.get_device_capability() print(f"Compute capability: {major}.{minor}") print(f"HAS_FA3: {HAS_FA3}") print() pytest.main([__file__, "-v", "-s"]) ================================================ FILE: tests/test_engine.py ================================================ """ Test Engine class. Example run: python -m pytest tests/test_engine.py -v """ import torch from nanochat.engine import KVCache, Engine from dataclasses import dataclass # ----------------------------------------------------------------------------- # Mock classes for testing Engine without loading a real model @dataclass class MockConfig: """Minimal config for Engine tests.""" n_kv_head: int = 4 n_head: int = 4 n_embd: int = 64 n_layer: int = 2 sequence_len: int = 128 class MockModel: """ Mock model that returns uniform logits over the vocab. This ensures that with temperature > 0, different samples should (with very high probability) produce different tokens. """ def __init__(self, vocab_size=262): # 256 bytes + 6 special tokens self.vocab_size = vocab_size self.config = MockConfig() self._device = torch.device("cpu") def get_device(self): return self._device def forward(self, ids, kv_cache=None): """Return uniform logits so sampling is spread across vocab.""" B, T = ids.shape # With FA3, flash_attn_with_kvcache updates cache in-place and we advance position if kv_cache is not None: kv_cache.advance(T) # Uniform logits -> equal probability for all tokens logits = torch.zeros(B, T, self.vocab_size) return logits class ByteTokenizer: """ Simple byte-level tokenizer for testing. Tokens 0-255 are raw bytes, 256+ are special tokens. """ def __init__(self): # Special tokens start at 256 self._special_tokens = { "<|python_start|>": 256, "<|python_end|>": 257, "<|output_start|>": 258, "<|output_end|>": 259, "<|assistant_end|>": 260, "<|bos|>": 261, } self._bos = 261 def encode_special(self, s): return self._special_tokens[s] def get_bos_token_id(self): return self._bos def encode(self, s, prepend=None): tokens = list(s.encode("utf-8")) # bytes 0-255 if prepend is not None: tokens = [prepend] + tokens return tokens def decode(self, tokens): # Filter out special tokens before decoding byte_tokens = [t for t in tokens if t < 256] return bytes(byte_tokens).decode("utf-8", errors="replace") def test_kv_cache_basic(): """Test basic KVCache functionality for FA3.""" batch_size = 2 num_heads = 3 seq_len = 64 head_dim = 5 num_layers = 6 kv_cache = KVCache( batch_size=batch_size, num_heads=num_heads, seq_len=seq_len, head_dim=head_dim, num_layers=num_layers, device="cpu", dtype=torch.float32, ) # Check initial state assert kv_cache.get_pos() == 0 assert kv_cache.k_cache.shape == (num_layers, batch_size, seq_len, num_heads, head_dim) assert kv_cache.v_cache.shape == (num_layers, batch_size, seq_len, num_heads, head_dim) # Test advance kv_cache.advance(10) assert kv_cache.get_pos() == 10 kv_cache.advance(5) assert kv_cache.get_pos() == 15 # Test reset kv_cache.reset() assert kv_cache.get_pos() == 0 # Test get_layer_cache returns correct views k_layer0, v_layer0 = kv_cache.get_layer_cache(0) assert k_layer0.shape == (batch_size, seq_len, num_heads, head_dim) assert v_layer0.shape == (batch_size, seq_len, num_heads, head_dim) def test_kv_cache_prefill(): """Test KVCache.prefill() copies data correctly.""" batch_size = 1 num_heads = 4 head_dim = 8 num_layers = 2 # Create source cache and advance it src_cache = KVCache( batch_size=batch_size, num_heads=num_heads, seq_len=32, head_dim=head_dim, num_layers=num_layers, device="cpu", dtype=torch.float32, ) # Write some data to source cache src_cache.k_cache[0, 0, :16, :, :] = 1.0 src_cache.v_cache[0, 0, :16, :, :] = 2.0 src_cache.advance(16) # Create destination cache with larger seq_len dst_cache = KVCache( batch_size=batch_size, num_heads=num_heads, seq_len=64, head_dim=head_dim, num_layers=num_layers, device="cpu", dtype=torch.float32, ) # Prefill dst_cache.prefill(src_cache) # Check position was copied assert dst_cache.get_pos() == 16 # Check data was copied assert (dst_cache.k_cache[0, 0, :16, :, :] == 1.0).all() assert (dst_cache.v_cache[0, 0, :16, :, :] == 2.0).all() def test_multi_sample_first_token_diversity(): """ Test that when generating multiple samples, each sample gets an independently sampled first token (not a broadcast of the same token to all rows). Previously, the first token after prefill was sampled once and broadcast to all rows, causing all samples to start identically. The fix expands the prefill logits to num_samples and samples independently for each row. With uniform logits over 262 tokens and 16 samples, the probability that all samples independently pick the same token is (1/262)^15 ≈ 10^-36. So if they're all identical, it indicates tokens are being broadcast instead of independently sampled. """ model = MockModel(vocab_size=262) tokenizer = ByteTokenizer() engine = Engine(model, tokenizer) # Generate 16 samples with temperature=1.0 (stochastic sampling) prompt_tokens = [261, 72, 101, 108, 108, 111] # + "Hello" num_samples = 16 # Collect the first generated token from each sample first_tokens = [] gen = engine.generate( prompt_tokens, num_samples=num_samples, max_tokens=1, # We only need the first token temperature=1.0, seed=42, ) for token_column, token_masks in gen: first_tokens = token_column # This is the first (and only) yield # With uniform distribution and 16 samples, they should NOT all be identical # If they are all identical, the bug exists (broadcasting instead of sampling) unique_tokens = set(first_tokens) assert len(unique_tokens) > 1, ( f"All {num_samples} samples got the same first token ({first_tokens[0]}). " f"With uniform logits, this is statistically impossible (~10^-36 probability) " f"unless tokens are being broadcast instead of independently sampled." ) def test_seed_reproducibility(): """Same seed must produce identical output.""" model = MockModel() engine = Engine(model, ByteTokenizer()) prompt = [261, 72, 101, 108, 108, 111] # + "Hello" for seed in [1, 42, 123, 999]: r1, _ = engine.generate_batch(prompt, max_tokens=5, seed=seed) r2, _ = engine.generate_batch(prompt, max_tokens=5, seed=seed) r3, _ = engine.generate_batch(prompt, max_tokens=5, seed=seed) assert r1 == r2 == r3, "Same seed must produce identical output for the same prompt." def test_temperature_zero_determinism(): """Temperature=0 is deterministic regardless of seed.""" model = MockModel() engine = Engine(model, ByteTokenizer()) prompt = [261, 72, 101, 108, 108, 111] r1, _ = engine.generate_batch(prompt, temperature=0.0, max_tokens=5, seed=1) r2, _ = engine.generate_batch(prompt, temperature=0.0, max_tokens=5, seed=42) r3, _ = engine.generate_batch(prompt, temperature=0.0, max_tokens=5, seed=123) assert r1 == r2 == r3, "Temperature=0 must result in the same output for the same prompt regardless of seed." def test_max_tokens_respected(): """Generation stops at max_tokens limit.""" model = MockModel() engine = Engine(model, ByteTokenizer()) prompt = [261, 72, 101, 108, 108, 111] for max_tokens in [1, 4, 16, 64]: results, _ = engine.generate_batch(prompt, max_tokens=max_tokens) num_generated_tokens = len(results[0]) - len(prompt) assert num_generated_tokens <= max_tokens, f"Generated {num_generated_tokens} tokens, expected max_tokens={max_tokens} or less." def test_num_samples_count(): """num_samples=N produces exactly N sequences.""" model = MockModel() engine = Engine(model, ByteTokenizer()) prompt = [261, 72, 101, 108, 108, 111] for num_samples in [1, 4, 16, 64]: results, _ = engine.generate_batch(prompt, num_samples=num_samples, max_tokens=3) assert len(results) == num_samples, f"Expected {num_samples} sequences from {num_samples} samples, got {len(results)}" def test_different_seeds_introduce_variation_when_temperature_nonzero(): """With temperature > 0, different seeds should introduce sampling variation.""" model = MockModel() engine = Engine(model, ByteTokenizer()) prompt = [261, 72, 101, 108, 108, 111] # + "Hello" outputs = set() for seed in [1, 42, 123, 999, 1000, 1001, 1002, 1003, 1004, 1005]: results, _ = engine.generate_batch( prompt, temperature=1.0, max_tokens=5, seed=seed, ) outputs.add(tuple(results[0])) # Sanity check: sampling actually introduces variation assert len(outputs) > 1, "All seeds produced the same output which is statistically highly improbable."