From 255f8b9af6d308a199601463b0062d618090a346 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 15 Jan 2026 23:30:11 +0000 Subject: [PATCH 01/33] cleanly separate cpu and gpu sections --- scripts/base_train.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/scripts/base_train.py b/scripts/base_train.py index bf4b8cf6..a9ee1c30 100644 --- a/scripts/base_train.py +++ b/scripts/base_train.py @@ -370,14 +370,15 @@ while True: for opt in optimizers: opt.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 + # 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.item() # EMA the training loss + 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(args.total_batch_size / dt) From 22a71aa3d30b50d6e7659ca0e8bcad9f3b7bfd98 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 15 Jan 2026 23:30:44 +0000 Subject: [PATCH 02/33] fuse adamw into a single torch compiled kernel similar to muon. it's about 1.7X faster, but overall it's so tiny that it's not making a major dent --- nanochat/adamw.py | 90 ++++++++++++++++++++++++++++++++++++----------- 1 file changed, 69 insertions(+), 21 deletions(-) diff --git a/nanochat/adamw.py b/nanochat/adamw.py index 48945b38..70ccf7b5 100644 --- a/nanochat/adamw.py +++ b/nanochat/adamw.py @@ -1,11 +1,42 @@ """ -Borrowed from modded-nanogpt. By Keller, @vagrawal, et al. -Not a general optimizer! But works for our specific use. +Distributed AdamW optimizer with a fused step function. +A bunch of ideas (e.g. dist comms in slices) are borrowed from modded-nanogpt. """ import torch import torch.distributed as dist from torch import Tensor +@torch.compile(dynamic=False, fullgraph=True) +def adamw_step_fused( + p: Tensor, + grad: Tensor, + exp_avg: Tensor, + exp_avg_sq: Tensor, + step_t: Tensor, + lr_t: Tensor, + beta1_t: Tensor, + beta2_t: Tensor, + eps_t: Tensor, + wd_t: Tensor, +) -> 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) + class DistAdamW(torch.optim.Optimizer): """ @@ -14,7 +45,26 @@ class DistAdamW(torch.optim.Optimizer): """ def __init__(self, param_groups, lr: float = 1e-3, betas: tuple[float, float] = (0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.01): defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay) + rank = dist.get_rank() + world_size = dist.get_world_size() + # Validate + if rank == 0: + for group in param_groups: + assert isinstance(group, dict), "expecting param_groups to be a list of dicts" + assert isinstance(group['params'], list), "expecting group['params'] to be a list of tensors" + for p in group['params']: + sliced = p.numel() >= 1024 + print(f"AdamW: 1 param of shape {p.shape}, sliced={sliced}") + if sliced: # large parameter tensors will be operated on in slices + assert p.shape[0] % world_size == 0, f"First dim of parameter shape {p.shape} must be divisible by world size {world_size}" super().__init__(param_groups, defaults) + # 0-D CPU tensors to avoid torch.compile recompilation when values change + self._step_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") + self._lr_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") + self._beta1_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") + self._beta2_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") + self._eps_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") + self._wd_t = torch.tensor(0.0, dtype=torch.float32, device="cpu") @torch.no_grad() def step(self): @@ -36,8 +86,7 @@ class DistAdamW(torch.optim.Optimizer): grad_slices.append(grad) else: is_small.append(False) - assert p.shape[0] % world_size == 0, f"First dim of parameter shape {p.shape} must be divisible by world size {world_size}" - rank_size = grad.shape[0] // world_size + rank_size = grad.shape[0] // world_size # p.shape[0] % world_size == 0 is checked in __init__ grad_slice = torch.empty_like(grad[:rank_size]) reduce_futures.append(dist.reduce_scatter_tensor(grad_slice, grad, op=dist.ReduceOp.AVG, async_op=True).get_future()) grad_slices.append(grad_slice) @@ -63,28 +112,27 @@ class DistAdamW(torch.optim.Optimizer): # State init if not state: - state['step'] = torch.tensor(0, dtype=torch.int64, device=p.device) + state['step'] = 0 state['exp_avg'] = torch.zeros_like(p_slice) state['exp_avg_sq'] = torch.zeros_like(p_slice) exp_avg = state['exp_avg'] exp_avg_sq = state['exp_avg_sq'] state['step'] += 1 - t = state['step'] - # weight decay - if wd != 0: - eff_weight_decay = lr * wd * getattr(p, "wd_mul", 1.0) - p_slice.mul_(1 - eff_weight_decay) - # update running averages - exp_avg.mul_(beta1).add_(g_slice, alpha=1 - beta1) - exp_avg_sq.mul_(beta2).addcmul_(g_slice, g_slice, value=1 - beta2) - # bias corrections - bias1 = 1 - beta1 ** t - bias2 = 1 - beta2 ** t - # compute step - denom = (exp_avg_sq / bias2).sqrt().add_(eps) - step_size = lr / bias1 - update = exp_avg.div(denom).mul_(step_size) - p_slice.add_(other=update, alpha=-1.0) + + # Fill 0-D tensors with current values + eff_wd = wd * getattr(p, "wd_mul", 1.0) + self._step_t.fill_(state['step']) + self._lr_t.fill_(lr) + self._beta1_t.fill_(beta1) + self._beta2_t.fill_(beta2) + self._eps_t.fill_(eps) + self._wd_t.fill_(eff_wd) + + # Fused update: weight_decay -> momentum -> bias_correction -> param_update + adamw_step_fused( + p_slice, g_slice, exp_avg, exp_avg_sq, + self._step_t, self._lr_t, self._beta1_t, self._beta2_t, self._eps_t, self._wd_t, + ) # Only large params need all_gather if not is_small[idx]: From bdcc030ffa97c829bd2e2c5841c776e1fb717aa2 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Thu, 15 Jan 2026 23:32:20 +0000 Subject: [PATCH 03/33] oops legacy spurious line now --- scripts/base_train.py | 1 - 1 file changed, 1 deletion(-) diff --git a/scripts/base_train.py b/scripts/base_train.py index a9ee1c30..5293cd88 100644 --- a/scripts/base_train.py +++ b/scripts/base_train.py @@ -208,7 +208,6 @@ if resuming: # ----------------------------------------------------------------------------- # Initialize the DataLoaders for train/val -tokens_dir = os.path.join(base_dir, "tokenized_data") 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) From d4ea28d4e2efc0b0491016721d01ec1afe83697e Mon Sep 17 00:00:00 2001 From: Sofie Van Landeghem Date: Fri, 16 Jan 2026 01:26:38 +0100 Subject: [PATCH 04/33] Fix args in readme (#438) * fix commands in readme, using new arg format * fix typo * add required -i flag to chat_eval example runs --- README.md | 4 ++-- scripts/chat_eval.py | 4 ++-- tasks/customjson.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index acb91110..9de2884c 100644 --- a/README.md +++ b/README.md @@ -82,10 +82,10 @@ That said, to give a sense, the example changes needed for the [speedrun.sh](spe python -m nanochat.dataset -n 450 & ... # use --depth to increase model size. to not oom, halve device batch size 32 -> 16: -torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- --depth=26 --device_batch_size=16 +torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- --depth=26 --device-batch-size=16 ... # make sure to use the same later during midtraining: -torchrun --standalone --nproc_per_node=8 -m scripts.mid_train -- --device_batch_size=16 +torchrun --standalone --nproc_per_node=8 -m scripts.mid_train -- --device-batch-size=16 ``` That's it! The biggest thing to pay attention to is making sure you have enough data shards to train on (the code will loop and do more epochs over the same training set otherwise, decreasing learning speed a bit), and managing your memory/VRAM, primarily by decreasing the `device_batch_size` until things fit (the scripts automatically compensate by increasing the number of gradient accumulation loops, simply turning parallel compute to sequential compute). diff --git a/scripts/chat_eval.py b/scripts/chat_eval.py index cae2f0f8..a5583035 100644 --- a/scripts/chat_eval.py +++ b/scripts/chat_eval.py @@ -4,8 +4,8 @@ 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 +python -m scripts.chat_eval -i mid -a ARC-Easy +torchrun --nproc_per_node=8 -m scripts.chat_eval -- -i mid -a ARC-Easy """ import argparse diff --git a/tasks/customjson.py b/tasks/customjson.py index e1b5f0b3..aeb1a3f7 100644 --- a/tasks/customjson.py +++ b/tasks/customjson.py @@ -25,7 +25,7 @@ class CustomJSON(Task): print("-" * 80) print(f"Warning: File {filepath} does not exist") print("HINT (Oct 21 2025)") - print("If you recently did a git pull and suddely see this, it might be due to the new addition of identity conversations") + print("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") From 7d1700c5215c1f10083976f6e0dc797e751501e2 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 00:40:59 +0000 Subject: [PATCH 05/33] add zstd lib --- pyproject.toml | 1 + uv.lock | 92 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 87a967f7..f3cd8d73 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "transformers>=4.57.3", "uvicorn>=0.36.0", "wandb>=0.21.3", + "zstandard>=0.25.0", ] [dependency-groups] diff --git a/uv.lock b/uv.lock index b168a2ff..dd766f89 100644 --- a/uv.lock +++ b/uv.lock @@ -1513,6 +1513,7 @@ dependencies = [ { name = "transformers" }, { name = "uvicorn" }, { name = "wandb" }, + { name = "zstandard" }, ] [package.optional-dependencies] @@ -1551,6 +1552,7 @@ requires-dist = [ { name = "transformers", specifier = ">=4.57.3" }, { name = "uvicorn", specifier = ">=0.36.0" }, { name = "wandb", specifier = ">=0.21.3" }, + { name = "zstandard", specifier = ">=0.25.0" }, ] provides-extras = ["cpu", "gpu"] @@ -3619,3 +3621,93 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/c3/b2e9f38bc3e11191981d57ea08cab2166e74ea770024a646617c9cddd9f6/yarl-1.20.1-cp313-cp313t-win_amd64.whl", hash = "sha256:541d050a355bbbc27e55d906bc91cb6fe42f96c01413dd0f4ed5a5240513874f", size = 93003, upload-time = "2025-06-10T00:45:27.752Z" }, { url = "https://files.pythonhosted.org/packages/b4/2d/2345fce04cfd4bee161bf1e7d9cdc702e3e16109021035dbb24db654a622/yarl-1.20.1-py3-none-any.whl", hash = "sha256:83b8eb083fe4683c6115795d9fc1cfaf2cbbefb19b3a1cb68f6527460f483a77", size = 46542, upload-time = "2025-06-10T00:46:07.521Z" }, ] + +[[package]] +name = "zstandard" +version = "0.25.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/7a/28efd1d371f1acd037ac64ed1c5e2b41514a6cc937dd6ab6a13ab9f0702f/zstandard-0.25.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e59fdc271772f6686e01e1b3b74537259800f57e24280be3f29c8a0deb1904dd", size = 795256, upload-time = "2025-09-14T22:15:56.415Z" }, + { url = "https://files.pythonhosted.org/packages/96/34/ef34ef77f1ee38fc8e4f9775217a613b452916e633c4f1d98f31db52c4a5/zstandard-0.25.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:4d441506e9b372386a5271c64125f72d5df6d2a8e8a2a45a0ae09b03cb781ef7", size = 640565, upload-time = "2025-09-14T22:15:58.177Z" }, + { url = "https://files.pythonhosted.org/packages/9d/1b/4fdb2c12eb58f31f28c4d28e8dc36611dd7205df8452e63f52fb6261d13e/zstandard-0.25.0-cp310-cp310-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:ab85470ab54c2cb96e176f40342d9ed41e58ca5733be6a893b730e7af9c40550", size = 5345306, upload-time = "2025-09-14T22:16:00.165Z" }, + { url = "https://files.pythonhosted.org/packages/73/28/a44bdece01bca027b079f0e00be3b6bd89a4df180071da59a3dd7381665b/zstandard-0.25.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e05ab82ea7753354bb054b92e2f288afb750e6b439ff6ca78af52939ebbc476d", size = 5055561, upload-time = "2025-09-14T22:16:02.22Z" }, + { url = "https://files.pythonhosted.org/packages/e9/74/68341185a4f32b274e0fc3410d5ad0750497e1acc20bd0f5b5f64ce17785/zstandard-0.25.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:78228d8a6a1c177a96b94f7e2e8d012c55f9c760761980da16ae7546a15a8e9b", size = 5402214, upload-time = "2025-09-14T22:16:04.109Z" }, + { url = "https://files.pythonhosted.org/packages/8b/67/f92e64e748fd6aaffe01e2b75a083c0c4fd27abe1c8747fee4555fcee7dd/zstandard-0.25.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:2b6bd67528ee8b5c5f10255735abc21aa106931f0dbaf297c7be0c886353c3d0", size = 5449703, upload-time = "2025-09-14T22:16:06.312Z" }, + { url = "https://files.pythonhosted.org/packages/fd/e5/6d36f92a197c3c17729a2125e29c169f460538a7d939a27eaaa6dcfcba8e/zstandard-0.25.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4b6d83057e713ff235a12e73916b6d356e3084fd3d14ced499d84240f3eecee0", size = 5556583, upload-time = "2025-09-14T22:16:08.457Z" }, + { url = "https://files.pythonhosted.org/packages/d7/83/41939e60d8d7ebfe2b747be022d0806953799140a702b90ffe214d557638/zstandard-0.25.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9174f4ed06f790a6869b41cba05b43eeb9a35f8993c4422ab853b705e8112bbd", size = 5045332, upload-time = "2025-09-14T22:16:10.444Z" }, + { url = "https://files.pythonhosted.org/packages/b3/87/d3ee185e3d1aa0133399893697ae91f221fda79deb61adbe998a7235c43f/zstandard-0.25.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:25f8f3cd45087d089aef5ba3848cd9efe3ad41163d3400862fb42f81a3a46701", size = 5572283, upload-time = "2025-09-14T22:16:12.128Z" }, + { url = "https://files.pythonhosted.org/packages/0a/1d/58635ae6104df96671076ac7d4ae7816838ce7debd94aecf83e30b7121b0/zstandard-0.25.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:3756b3e9da9b83da1796f8809dd57cb024f838b9eeafde28f3cb472012797ac1", size = 4959754, upload-time = "2025-09-14T22:16:14.225Z" }, + { url = "https://files.pythonhosted.org/packages/75/d6/57e9cb0a9983e9a229dd8fd2e6e96593ef2aa82a3907188436f22b111ccd/zstandard-0.25.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:81dad8d145d8fd981b2962b686b2241d3a1ea07733e76a2f15435dfb7fb60150", size = 5266477, upload-time = "2025-09-14T22:16:16.343Z" }, + { url = "https://files.pythonhosted.org/packages/d1/a9/ee891e5edf33a6ebce0a028726f0bbd8567effe20fe3d5808c42323e8542/zstandard-0.25.0-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:a5a419712cf88862a45a23def0ae063686db3d324cec7edbe40509d1a79a0aab", size = 5440914, upload-time = "2025-09-14T22:16:18.453Z" }, + { url = "https://files.pythonhosted.org/packages/58/08/a8522c28c08031a9521f27abc6f78dbdee7312a7463dd2cfc658b813323b/zstandard-0.25.0-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:e7360eae90809efd19b886e59a09dad07da4ca9ba096752e61a2e03c8aca188e", size = 5819847, upload-time = "2025-09-14T22:16:20.559Z" }, + { url = "https://files.pythonhosted.org/packages/6f/11/4c91411805c3f7b6f31c60e78ce347ca48f6f16d552fc659af6ec3b73202/zstandard-0.25.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:75ffc32a569fb049499e63ce68c743155477610532da1eb38e7f24bf7cd29e74", size = 5363131, upload-time = "2025-09-14T22:16:22.206Z" }, + { url = "https://files.pythonhosted.org/packages/ef/d6/8c4bd38a3b24c4c7676a7a3d8de85d6ee7a983602a734b9f9cdefb04a5d6/zstandard-0.25.0-cp310-cp310-win32.whl", hash = "sha256:106281ae350e494f4ac8a80470e66d1fe27e497052c8d9c3b95dc4cf1ade81aa", size = 436469, upload-time = "2025-09-14T22:16:25.002Z" }, + { url = "https://files.pythonhosted.org/packages/93/90/96d50ad417a8ace5f841b3228e93d1bb13e6ad356737f42e2dde30d8bd68/zstandard-0.25.0-cp310-cp310-win_amd64.whl", hash = "sha256:ea9d54cc3d8064260114a0bbf3479fc4a98b21dffc89b3459edd506b69262f6e", size = 506100, upload-time = "2025-09-14T22:16:23.569Z" }, + { url = "https://files.pythonhosted.org/packages/2a/83/c3ca27c363d104980f1c9cee1101cc8ba724ac8c28a033ede6aab89585b1/zstandard-0.25.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:933b65d7680ea337180733cf9e87293cc5500cc0eb3fc8769f4d3c88d724ec5c", size = 795254, upload-time = "2025-09-14T22:16:26.137Z" }, + { url = "https://files.pythonhosted.org/packages/ac/4d/e66465c5411a7cf4866aeadc7d108081d8ceba9bc7abe6b14aa21c671ec3/zstandard-0.25.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a3f79487c687b1fc69f19e487cd949bf3aae653d181dfb5fde3bf6d18894706f", size = 640559, upload-time = "2025-09-14T22:16:27.973Z" }, + { url = "https://files.pythonhosted.org/packages/12/56/354fe655905f290d3b147b33fe946b0f27e791e4b50a5f004c802cb3eb7b/zstandard-0.25.0-cp311-cp311-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:0bbc9a0c65ce0eea3c34a691e3c4b6889f5f3909ba4822ab385fab9057099431", size = 5348020, upload-time = "2025-09-14T22:16:29.523Z" }, + { url = "https://files.pythonhosted.org/packages/3b/13/2b7ed68bd85e69a2069bcc72141d378f22cae5a0f3b353a2c8f50ef30c1b/zstandard-0.25.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:01582723b3ccd6939ab7b3a78622c573799d5d8737b534b86d0e06ac18dbde4a", size = 5058126, upload-time = "2025-09-14T22:16:31.811Z" }, + { url = "https://files.pythonhosted.org/packages/c9/dd/fdaf0674f4b10d92cb120ccff58bbb6626bf8368f00ebfd2a41ba4a0dc99/zstandard-0.25.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:5f1ad7bf88535edcf30038f6919abe087f606f62c00a87d7e33e7fc57cb69fcc", size = 5405390, upload-time = "2025-09-14T22:16:33.486Z" }, + { url = "https://files.pythonhosted.org/packages/0f/67/354d1555575bc2490435f90d67ca4dd65238ff2f119f30f72d5cde09c2ad/zstandard-0.25.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:06acb75eebeedb77b69048031282737717a63e71e4ae3f77cc0c3b9508320df6", size = 5452914, upload-time = "2025-09-14T22:16:35.277Z" }, + { url = "https://files.pythonhosted.org/packages/bb/1f/e9cfd801a3f9190bf3e759c422bbfd2247db9d7f3d54a56ecde70137791a/zstandard-0.25.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9300d02ea7c6506f00e627e287e0492a5eb0371ec1670ae852fefffa6164b072", size = 5559635, upload-time = "2025-09-14T22:16:37.141Z" }, + { url = "https://files.pythonhosted.org/packages/21/88/5ba550f797ca953a52d708c8e4f380959e7e3280af029e38fbf47b55916e/zstandard-0.25.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:bfd06b1c5584b657a2892a6014c2f4c20e0db0208c159148fa78c65f7e0b0277", size = 5048277, upload-time = "2025-09-14T22:16:38.807Z" }, + { url = "https://files.pythonhosted.org/packages/46/c0/ca3e533b4fa03112facbe7fbe7779cb1ebec215688e5df576fe5429172e0/zstandard-0.25.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:f373da2c1757bb7f1acaf09369cdc1d51d84131e50d5fa9863982fd626466313", size = 5574377, upload-time = "2025-09-14T22:16:40.523Z" }, + { url = "https://files.pythonhosted.org/packages/12/9b/3fb626390113f272abd0799fd677ea33d5fc3ec185e62e6be534493c4b60/zstandard-0.25.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:6c0e5a65158a7946e7a7affa6418878ef97ab66636f13353b8502d7ea03c8097", size = 4961493, upload-time = "2025-09-14T22:16:43.3Z" }, + { url = "https://files.pythonhosted.org/packages/cb/d3/23094a6b6a4b1343b27ae68249daa17ae0651fcfec9ed4de09d14b940285/zstandard-0.25.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:c8e167d5adf59476fa3e37bee730890e389410c354771a62e3c076c86f9f7778", size = 5269018, upload-time = "2025-09-14T22:16:45.292Z" }, + { url = "https://files.pythonhosted.org/packages/8c/a7/bb5a0c1c0f3f4b5e9d5b55198e39de91e04ba7c205cc46fcb0f95f0383c1/zstandard-0.25.0-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:98750a309eb2f020da61e727de7d7ba3c57c97cf6213f6f6277bb7fb42a8e065", size = 5443672, upload-time = "2025-09-14T22:16:47.076Z" }, + { url = "https://files.pythonhosted.org/packages/27/22/503347aa08d073993f25109c36c8d9f029c7d5949198050962cb568dfa5e/zstandard-0.25.0-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:22a086cff1b6ceca18a8dd6096ec631e430e93a8e70a9ca5efa7561a00f826fa", size = 5822753, upload-time = "2025-09-14T22:16:49.316Z" }, + { url = "https://files.pythonhosted.org/packages/e2/be/94267dc6ee64f0f8ba2b2ae7c7a2df934a816baaa7291db9e1aa77394c3c/zstandard-0.25.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:72d35d7aa0bba323965da807a462b0966c91608ef3a48ba761678cb20ce5d8b7", size = 5366047, upload-time = "2025-09-14T22:16:51.328Z" }, + { url = "https://files.pythonhosted.org/packages/7b/a3/732893eab0a3a7aecff8b99052fecf9f605cf0fb5fb6d0290e36beee47a4/zstandard-0.25.0-cp311-cp311-win32.whl", hash = "sha256:f5aeea11ded7320a84dcdd62a3d95b5186834224a9e55b92ccae35d21a8b63d4", size = 436484, upload-time = "2025-09-14T22:16:55.005Z" }, + { url = "https://files.pythonhosted.org/packages/43/a3/c6155f5c1cce691cb80dfd38627046e50af3ee9ddc5d0b45b9b063bfb8c9/zstandard-0.25.0-cp311-cp311-win_amd64.whl", hash = "sha256:daab68faadb847063d0c56f361a289c4f268706b598afbf9ad113cbe5c38b6b2", size = 506183, upload-time = "2025-09-14T22:16:52.753Z" }, + { url = "https://files.pythonhosted.org/packages/8c/3e/8945ab86a0820cc0e0cdbf38086a92868a9172020fdab8a03ac19662b0e5/zstandard-0.25.0-cp311-cp311-win_arm64.whl", hash = "sha256:22a06c5df3751bb7dc67406f5374734ccee8ed37fc5981bf1ad7041831fa1137", size = 462533, upload-time = "2025-09-14T22:16:53.878Z" }, + { url = "https://files.pythonhosted.org/packages/82/fc/f26eb6ef91ae723a03e16eddb198abcfce2bc5a42e224d44cc8b6765e57e/zstandard-0.25.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7b3c3a3ab9daa3eed242d6ecceead93aebbb8f5f84318d82cee643e019c4b73b", size = 795738, upload-time = "2025-09-14T22:16:56.237Z" }, + { url = "https://files.pythonhosted.org/packages/aa/1c/d920d64b22f8dd028a8b90e2d756e431a5d86194caa78e3819c7bf53b4b3/zstandard-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:913cbd31a400febff93b564a23e17c3ed2d56c064006f54efec210d586171c00", size = 640436, upload-time = "2025-09-14T22:16:57.774Z" }, + { url = "https://files.pythonhosted.org/packages/53/6c/288c3f0bd9fcfe9ca41e2c2fbfd17b2097f6af57b62a81161941f09afa76/zstandard-0.25.0-cp312-cp312-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:011d388c76b11a0c165374ce660ce2c8efa8e5d87f34996aa80f9c0816698b64", size = 5343019, upload-time = "2025-09-14T22:16:59.302Z" }, + { url = "https://files.pythonhosted.org/packages/1e/15/efef5a2f204a64bdb5571e6161d49f7ef0fffdbca953a615efbec045f60f/zstandard-0.25.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dffecc361d079bb48d7caef5d673c88c8988d3d33fb74ab95b7ee6da42652ea", size = 5063012, upload-time = "2025-09-14T22:17:01.156Z" }, + { url = "https://files.pythonhosted.org/packages/b7/37/a6ce629ffdb43959e92e87ebdaeebb5ac81c944b6a75c9c47e300f85abdf/zstandard-0.25.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:7149623bba7fdf7e7f24312953bcf73cae103db8cae49f8154dd1eadc8a29ecb", size = 5394148, upload-time = "2025-09-14T22:17:03.091Z" }, + { url = "https://files.pythonhosted.org/packages/e3/79/2bf870b3abeb5c070fe2d670a5a8d1057a8270f125ef7676d29ea900f496/zstandard-0.25.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:6a573a35693e03cf1d67799fd01b50ff578515a8aeadd4595d2a7fa9f3ec002a", size = 5451652, upload-time = "2025-09-14T22:17:04.979Z" }, + { url = "https://files.pythonhosted.org/packages/53/60/7be26e610767316c028a2cbedb9a3beabdbe33e2182c373f71a1c0b88f36/zstandard-0.25.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5a56ba0db2d244117ed744dfa8f6f5b366e14148e00de44723413b2f3938a902", size = 5546993, upload-time = "2025-09-14T22:17:06.781Z" }, + { url = "https://files.pythonhosted.org/packages/85/c7/3483ad9ff0662623f3648479b0380d2de5510abf00990468c286c6b04017/zstandard-0.25.0-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:10ef2a79ab8e2974e2075fb984e5b9806c64134810fac21576f0668e7ea19f8f", size = 5046806, upload-time = "2025-09-14T22:17:08.415Z" }, + { url = "https://files.pythonhosted.org/packages/08/b3/206883dd25b8d1591a1caa44b54c2aad84badccf2f1de9e2d60a446f9a25/zstandard-0.25.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:aaf21ba8fb76d102b696781bddaa0954b782536446083ae3fdaa6f16b25a1c4b", size = 5576659, upload-time = "2025-09-14T22:17:10.164Z" }, + { url = "https://files.pythonhosted.org/packages/9d/31/76c0779101453e6c117b0ff22565865c54f48f8bd807df2b00c2c404b8e0/zstandard-0.25.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1869da9571d5e94a85a5e8d57e4e8807b175c9e4a6294e3b66fa4efb074d90f6", size = 4953933, upload-time = "2025-09-14T22:17:11.857Z" }, + { url = "https://files.pythonhosted.org/packages/18/e1/97680c664a1bf9a247a280a053d98e251424af51f1b196c6d52f117c9720/zstandard-0.25.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:809c5bcb2c67cd0ed81e9229d227d4ca28f82d0f778fc5fea624a9def3963f91", size = 5268008, upload-time = "2025-09-14T22:17:13.627Z" }, + { url = "https://files.pythonhosted.org/packages/1e/73/316e4010de585ac798e154e88fd81bb16afc5c5cb1a72eeb16dd37e8024a/zstandard-0.25.0-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:f27662e4f7dbf9f9c12391cb37b4c4c3cb90ffbd3b1fb9284dadbbb8935fa708", size = 5433517, upload-time = "2025-09-14T22:17:16.103Z" }, + { url = "https://files.pythonhosted.org/packages/5b/60/dd0f8cfa8129c5a0ce3ea6b7f70be5b33d2618013a161e1ff26c2b39787c/zstandard-0.25.0-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:99c0c846e6e61718715a3c9437ccc625de26593fea60189567f0118dc9db7512", size = 5814292, upload-time = "2025-09-14T22:17:17.827Z" }, + { url = "https://files.pythonhosted.org/packages/fc/5f/75aafd4b9d11b5407b641b8e41a57864097663699f23e9ad4dbb91dc6bfe/zstandard-0.25.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:474d2596a2dbc241a556e965fb76002c1ce655445e4e3bf38e5477d413165ffa", size = 5360237, upload-time = "2025-09-14T22:17:19.954Z" }, + { url = "https://files.pythonhosted.org/packages/ff/8d/0309daffea4fcac7981021dbf21cdb2e3427a9e76bafbcdbdf5392ff99a4/zstandard-0.25.0-cp312-cp312-win32.whl", hash = "sha256:23ebc8f17a03133b4426bcc04aabd68f8236eb78c3760f12783385171b0fd8bd", size = 436922, upload-time = "2025-09-14T22:17:24.398Z" }, + { url = "https://files.pythonhosted.org/packages/79/3b/fa54d9015f945330510cb5d0b0501e8253c127cca7ebe8ba46a965df18c5/zstandard-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:ffef5a74088f1e09947aecf91011136665152e0b4b359c42be3373897fb39b01", size = 506276, upload-time = "2025-09-14T22:17:21.429Z" }, + { url = "https://files.pythonhosted.org/packages/ea/6b/8b51697e5319b1f9ac71087b0af9a40d8a6288ff8025c36486e0c12abcc4/zstandard-0.25.0-cp312-cp312-win_arm64.whl", hash = "sha256:181eb40e0b6a29b3cd2849f825e0fa34397f649170673d385f3598ae17cca2e9", size = 462679, upload-time = "2025-09-14T22:17:23.147Z" }, + { url = "https://files.pythonhosted.org/packages/35/0b/8df9c4ad06af91d39e94fa96cc010a24ac4ef1378d3efab9223cc8593d40/zstandard-0.25.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ec996f12524f88e151c339688c3897194821d7f03081ab35d31d1e12ec975e94", size = 795735, upload-time = "2025-09-14T22:17:26.042Z" }, + { url = "https://files.pythonhosted.org/packages/3f/06/9ae96a3e5dcfd119377ba33d4c42a7d89da1efabd5cb3e366b156c45ff4d/zstandard-0.25.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a1a4ae2dec3993a32247995bdfe367fc3266da832d82f8438c8570f989753de1", size = 640440, upload-time = "2025-09-14T22:17:27.366Z" }, + { url = "https://files.pythonhosted.org/packages/d9/14/933d27204c2bd404229c69f445862454dcc101cd69ef8c6068f15aaec12c/zstandard-0.25.0-cp313-cp313-manylinux2010_i686.manylinux2014_i686.manylinux_2_12_i686.manylinux_2_17_i686.whl", hash = "sha256:e96594a5537722fdfb79951672a2a63aec5ebfb823e7560586f7484819f2a08f", size = 5343070, upload-time = "2025-09-14T22:17:28.896Z" }, + { url = "https://files.pythonhosted.org/packages/6d/db/ddb11011826ed7db9d0e485d13df79b58586bfdec56e5c84a928a9a78c1c/zstandard-0.25.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bfc4e20784722098822e3eee42b8e576b379ed72cca4a7cb856ae733e62192ea", size = 5063001, upload-time = "2025-09-14T22:17:31.044Z" }, + { url = "https://files.pythonhosted.org/packages/db/00/87466ea3f99599d02a5238498b87bf84a6348290c19571051839ca943777/zstandard-0.25.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:457ed498fc58cdc12fc48f7950e02740d4f7ae9493dd4ab2168a47c93c31298e", size = 5394120, upload-time = "2025-09-14T22:17:32.711Z" }, + { url = "https://files.pythonhosted.org/packages/2b/95/fc5531d9c618a679a20ff6c29e2b3ef1d1f4ad66c5e161ae6ff847d102a9/zstandard-0.25.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:fd7a5004eb1980d3cefe26b2685bcb0b17989901a70a1040d1ac86f1d898c551", size = 5451230, upload-time = "2025-09-14T22:17:34.41Z" }, + { url = "https://files.pythonhosted.org/packages/63/4b/e3678b4e776db00f9f7b2fe58e547e8928ef32727d7a1ff01dea010f3f13/zstandard-0.25.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8e735494da3db08694d26480f1493ad2cf86e99bdd53e8e9771b2752a5c0246a", size = 5547173, upload-time = "2025-09-14T22:17:36.084Z" }, + { url = "https://files.pythonhosted.org/packages/4e/d5/ba05ed95c6b8ec30bd468dfeab20589f2cf709b5c940483e31d991f2ca58/zstandard-0.25.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3a39c94ad7866160a4a46d772e43311a743c316942037671beb264e395bdd611", size = 5046736, upload-time = "2025-09-14T22:17:37.891Z" }, + { url = "https://files.pythonhosted.org/packages/50/d5/870aa06b3a76c73eced65c044b92286a3c4e00554005ff51962deef28e28/zstandard-0.25.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:172de1f06947577d3a3005416977cce6168f2261284c02080e7ad0185faeced3", size = 5576368, upload-time = "2025-09-14T22:17:40.206Z" }, + { url = "https://files.pythonhosted.org/packages/5d/35/398dc2ffc89d304d59bc12f0fdd931b4ce455bddf7038a0a67733a25f550/zstandard-0.25.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3c83b0188c852a47cd13ef3bf9209fb0a77fa5374958b8c53aaa699398c6bd7b", size = 4954022, upload-time = "2025-09-14T22:17:41.879Z" }, + { url = "https://files.pythonhosted.org/packages/9a/5c/36ba1e5507d56d2213202ec2b05e8541734af5f2ce378c5d1ceaf4d88dc4/zstandard-0.25.0-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:1673b7199bbe763365b81a4f3252b8e80f44c9e323fc42940dc8843bfeaf9851", size = 5267889, upload-time = "2025-09-14T22:17:43.577Z" }, + { url = "https://files.pythonhosted.org/packages/70/e8/2ec6b6fb7358b2ec0113ae202647ca7c0e9d15b61c005ae5225ad0995df5/zstandard-0.25.0-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:0be7622c37c183406f3dbf0cba104118eb16a4ea7359eeb5752f0794882fc250", size = 5433952, upload-time = "2025-09-14T22:17:45.271Z" }, + { url = "https://files.pythonhosted.org/packages/7b/01/b5f4d4dbc59ef193e870495c6f1275f5b2928e01ff5a81fecb22a06e22fb/zstandard-0.25.0-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:5f5e4c2a23ca271c218ac025bd7d635597048b366d6f31f420aaeb715239fc98", size = 5814054, upload-time = "2025-09-14T22:17:47.08Z" }, + { url = "https://files.pythonhosted.org/packages/b2/e5/fbd822d5c6f427cf158316d012c5a12f233473c2f9c5fe5ab1ae5d21f3d8/zstandard-0.25.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4f187a0bb61b35119d1926aee039524d1f93aaf38a9916b8c4b78ac8514a0aaf", size = 5360113, upload-time = "2025-09-14T22:17:48.893Z" }, + { url = "https://files.pythonhosted.org/packages/8e/e0/69a553d2047f9a2c7347caa225bb3a63b6d7704ad74610cb7823baa08ed7/zstandard-0.25.0-cp313-cp313-win32.whl", hash = "sha256:7030defa83eef3e51ff26f0b7bfb229f0204b66fe18e04359ce3474ac33cbc09", size = 436936, upload-time = "2025-09-14T22:17:52.658Z" }, + { url = "https://files.pythonhosted.org/packages/d9/82/b9c06c870f3bd8767c201f1edbdf9e8dc34be5b0fbc5682c4f80fe948475/zstandard-0.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:1f830a0dac88719af0ae43b8b2d6aef487d437036468ef3c2ea59c51f9d55fd5", size = 506232, upload-time = "2025-09-14T22:17:50.402Z" }, + { url = "https://files.pythonhosted.org/packages/d4/57/60c3c01243bb81d381c9916e2a6d9e149ab8627c0c7d7abb2d73384b3c0c/zstandard-0.25.0-cp313-cp313-win_arm64.whl", hash = "sha256:85304a43f4d513f5464ceb938aa02c1e78c2943b29f44a750b48b25ac999a049", size = 462671, upload-time = "2025-09-14T22:17:51.533Z" }, + { url = "https://files.pythonhosted.org/packages/3d/5c/f8923b595b55fe49e30612987ad8bf053aef555c14f05bb659dd5dbe3e8a/zstandard-0.25.0-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:e29f0cf06974c899b2c188ef7f783607dbef36da4c242eb6c82dcd8b512855e3", size = 795887, upload-time = "2025-09-14T22:17:54.198Z" }, + { url = "https://files.pythonhosted.org/packages/8d/09/d0a2a14fc3439c5f874042dca72a79c70a532090b7ba0003be73fee37ae2/zstandard-0.25.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:05df5136bc5a011f33cd25bc9f506e7426c0c9b3f9954f056831ce68f3b6689f", size = 640658, upload-time = "2025-09-14T22:17:55.423Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7c/8b6b71b1ddd517f68ffb55e10834388d4f793c49c6b83effaaa05785b0b4/zstandard-0.25.0-cp314-cp314-manylinux2010_i686.manylinux_2_12_i686.manylinux_2_28_i686.whl", hash = "sha256:f604efd28f239cc21b3adb53eb061e2a205dc164be408e553b41ba2ffe0ca15c", size = 5379849, upload-time = "2025-09-14T22:17:57.372Z" }, + { url = "https://files.pythonhosted.org/packages/a4/86/a48e56320d0a17189ab7a42645387334fba2200e904ee47fc5a26c1fd8ca/zstandard-0.25.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:223415140608d0f0da010499eaa8ccdb9af210a543fac54bce15babbcfc78439", size = 5058095, upload-time = "2025-09-14T22:17:59.498Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ad/eb659984ee2c0a779f9d06dbfe45e2dc39d99ff40a319895df2d3d9a48e5/zstandard-0.25.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2e54296a283f3ab5a26fc9b8b5d4978ea0532f37b231644f367aa588930aa043", size = 5551751, upload-time = "2025-09-14T22:18:01.618Z" }, + { url = "https://files.pythonhosted.org/packages/61/b3/b637faea43677eb7bd42ab204dfb7053bd5c4582bfe6b1baefa80ac0c47b/zstandard-0.25.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ca54090275939dc8ec5dea2d2afb400e0f83444b2fc24e07df7fdef677110859", size = 6364818, upload-time = "2025-09-14T22:18:03.769Z" }, + { url = "https://files.pythonhosted.org/packages/31/dc/cc50210e11e465c975462439a492516a73300ab8caa8f5e0902544fd748b/zstandard-0.25.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e09bb6252b6476d8d56100e8147b803befa9a12cea144bbe629dd508800d1ad0", size = 5560402, upload-time = "2025-09-14T22:18:05.954Z" }, + { url = "https://files.pythonhosted.org/packages/c9/ae/56523ae9c142f0c08efd5e868a6da613ae76614eca1305259c3bf6a0ed43/zstandard-0.25.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a9ec8c642d1ec73287ae3e726792dd86c96f5681eb8df274a757bf62b750eae7", size = 4955108, upload-time = "2025-09-14T22:18:07.68Z" }, + { url = "https://files.pythonhosted.org/packages/98/cf/c899f2d6df0840d5e384cf4c4121458c72802e8bda19691f3b16619f51e9/zstandard-0.25.0-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:a4089a10e598eae6393756b036e0f419e8c1d60f44a831520f9af41c14216cf2", size = 5269248, upload-time = "2025-09-14T22:18:09.753Z" }, + { url = "https://files.pythonhosted.org/packages/1b/c0/59e912a531d91e1c192d3085fc0f6fb2852753c301a812d856d857ea03c6/zstandard-0.25.0-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:f67e8f1a324a900e75b5e28ffb152bcac9fbed1cc7b43f99cd90f395c4375344", size = 5430330, upload-time = "2025-09-14T22:18:11.966Z" }, + { url = "https://files.pythonhosted.org/packages/a0/1d/7e31db1240de2df22a58e2ea9a93fc6e38cc29353e660c0272b6735d6669/zstandard-0.25.0-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:9654dbc012d8b06fc3d19cc825af3f7bf8ae242226df5f83936cb39f5fdc846c", size = 5811123, upload-time = "2025-09-14T22:18:13.907Z" }, + { url = "https://files.pythonhosted.org/packages/f6/49/fac46df5ad353d50535e118d6983069df68ca5908d4d65b8c466150a4ff1/zstandard-0.25.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:4203ce3b31aec23012d3a4cf4a2ed64d12fea5269c49aed5e4c3611b938e4088", size = 5359591, upload-time = "2025-09-14T22:18:16.465Z" }, + { url = "https://files.pythonhosted.org/packages/c2/38/f249a2050ad1eea0bb364046153942e34abba95dd5520af199aed86fbb49/zstandard-0.25.0-cp314-cp314-win32.whl", hash = "sha256:da469dc041701583e34de852d8634703550348d5822e66a0c827d39b05365b12", size = 444513, upload-time = "2025-09-14T22:18:20.61Z" }, + { url = "https://files.pythonhosted.org/packages/3a/43/241f9615bcf8ba8903b3f0432da069e857fc4fd1783bd26183db53c4804b/zstandard-0.25.0-cp314-cp314-win_amd64.whl", hash = "sha256:c19bcdd826e95671065f8692b5a4aa95c52dc7a02a4c5a0cac46deb879a017a2", size = 516118, upload-time = "2025-09-14T22:18:17.849Z" }, + { url = "https://files.pythonhosted.org/packages/f0/ef/da163ce2450ed4febf6467d77ccb4cd52c4c30ab45624bad26ca0a27260c/zstandard-0.25.0-cp314-cp314-win_arm64.whl", hash = "sha256:d7541afd73985c630bafcd6338d2518ae96060075f9463d7dc14cfb33514383d", size = 476940, upload-time = "2025-09-14T22:18:19.088Z" }, +] From 747ed4491f7fe77b1f99a385309804a3f2cca353 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 00:43:54 +0000 Subject: [PATCH 06/33] add negative result on olmo3 pretraining mix --- dev/LOG.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/dev/LOG.md b/dev/LOG.md index 5f6e1d7f..2a94daaf 100644 --- a/dev/LOG.md +++ b/dev/LOG.md @@ -4,6 +4,14 @@ A running summary documenting some experiments and findings. Started ~Jan 7 2026 --- +## 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. From fbf2bbea25da6ada9d77db4c042937f59205823f Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 02:21:17 +0000 Subject: [PATCH 07/33] update log with a bunch of attempts --- dev/LOG.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/dev/LOG.md b/dev/LOG.md index 2a94daaf..d0dc1b18 100644 --- a/dev/LOG.md +++ b/dev/LOG.md @@ -4,6 +4,22 @@ A running summary documenting some experiments and findings. Started ~Jan 7 2026 --- +## 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. @@ -12,6 +28,8 @@ I am still looking to try the [DCLM dataset](https://arxiv.org/abs/2406.11794), 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. From 50413d2d67b31db8c4f1b593808f98eb23b4401a Mon Sep 17 00:00:00 2001 From: Haoyu Wang <32129905+why2011btv@users.noreply.github.com> Date: Fri, 16 Jan 2026 01:03:42 -0500 Subject: [PATCH 08/33] typo in comments: change "GAPO" to "DAPO" --- scripts/chat_rl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/chat_rl.py b/scripts/chat_rl.py index b0697f36..eb8e48e5 100644 --- a/scripts/chat_rl.py +++ b/scripts/chat_rl.py @@ -6,7 +6,7 @@ simpler and more similar to just REINFORCE: 1) Delete trust region, so there is no KL regularization to a reference model 2) We are on policy, so there's no need for PPO ratio+clip. -3) We use GAPO style normalization that is token-level, not sequence-level. +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: From 8203efa9190e9d8da7419de29599e91a9242b968 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 17:37:51 +0000 Subject: [PATCH 09/33] implement flash attention 3 fallback to pytorch sdpa by touching as few lines of code as possible in main files and keeping all implementation to a single file. add tests. add helpful warning messages for the user. --- nanochat/gpt.py | 12 +- scripts/base_train.py | 13 ++ tests/test_attention_fallback.py | 338 +++++++++++++++++++++++++++++++ 3 files changed, 354 insertions(+), 9 deletions(-) create mode 100644 tests/test_attention_fallback.py diff --git a/nanochat/gpt.py b/nanochat/gpt.py index 81ccb0ca..86f440bf 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -23,13 +23,8 @@ from nanochat.common import get_dist_info, print0 from nanochat.muon import Muon, DistMuon from nanochat.adamw import DistAdamW -# Load Flash Attention 3 from HuggingFace Hub (and silence the progress bar) -import os -os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" -# Official docs of FA3 label it as "beta" and want you to install FA3 from source, which is a pain. -# Wishing for official FA3 wheels soon, for now this seems to be a fast way to get them (ty varunneal) -from kernels import get_kernel -flash_attn = get_kernel('varunneal/flash-attention-3').flash_attn_interface +# 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: @@ -87,8 +82,7 @@ class CausalSelfAttention(nn.Module): q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) q, k = norm(q), norm(k) # QK norm - # Attention with Flash Attention 3 - # FA3 handles GQA automatically when n_kv_heads < n_heads + # 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 diff --git a/scripts/base_train.py b/scripts/base_train.py index 5293cd88..c61986e6 100644 --- a/scripts/base_train.py +++ b/scripts/base_train.py @@ -27,6 +27,7 @@ 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_model print_banner() @@ -86,6 +87,18 @@ get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else l 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 +if HAS_FA3: + print0("✓ Using Flash Attention 3 (Hopper GPU detected), efficient, new and awesome.") +else: + print0("!" * 80) + 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, also we need the vocab size tokenizer = get_tokenizer() token_bytes = get_token_bytes(device=device) diff --git a/tests/test_attention_fallback.py b/tests/test_attention_fallback.py new file mode 100644 index 00000000..2cf3ed77 --- /dev/null +++ b/tests/test_attention_fallback.py @@ -0,0 +1,338 @@ +""" +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).""" + fa_module._override_impl = impl + + +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_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"]) From b62a5bc44aafef02eba6c39236180aa424a82674 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 17:39:41 +0000 Subject: [PATCH 10/33] naturally i failed to include the actual code in the previous commit facepalm --- nanochat/flash_attention.py | 178 ++++++++++++++++++++++++++++++++++++ 1 file changed, 178 insertions(+) create mode 100644 nanochat/flash_attention.py diff --git a/nanochat/flash_attention.py b/nanochat/flash_attention.py new file mode 100644 index 00000000..5d27e5f9 --- /dev/null +++ b/nanochat/flash_attention.py @@ -0,0 +1,178 @@ +""" +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, 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).""" + if not torch.cuda.is_available(): + return None + try: + major, _ = torch.cuda.get_device_capability() + if major < 9: # Hopper is sm90 + 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 _use_fa3(): + """Determine whether to use FA3 based on availability and override.""" + if _override_impl == 'fa3': + assert HAS_FA3, "Cannot override to FA3: not available on this hardware" + return True + if _override_impl == 'sdpa': + return False + return HAS_FA3 # auto + + +# ============================================================================= +# 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: + return F.scaled_dot_product_attention(q, k, v, is_causal=False, enable_gqa=enable_gqa) + + # Need explicit mask + device = q.device + if Tq == Tk: + # Causal + sliding window + mask = torch.tril(torch.ones(Tq, Tk, device=device, dtype=torch.bool)) + if window > 0 and window < Tq: + row_idx = torch.arange(Tq, device=device).unsqueeze(1) + col_idx = torch.arange(Tk, device=device).unsqueeze(0) + mask = mask & ((row_idx - col_idx) <= window) + else: + # Chunk inference: attend to prefix + causal within chunk + prefix_len = Tk - Tq + mask = torch.zeros(Tq, Tk, device=device, dtype=torch.bool) + mask[:, :prefix_len] = True + mask[:, prefix_len:] = torch.tril(torch.ones(Tq, Tq, device=device, dtype=torch.bool)) + + 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, +) From 184d4c12b1d01b098aeffa021be5168e454afad0 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 18:25:04 +0000 Subject: [PATCH 11/33] also add to log about the FA3 changes --- dev/LOG.md | 44 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/dev/LOG.md b/dev/LOG.md index d0dc1b18..ae518c83 100644 --- a/dev/LOG.md +++ b/dev/LOG.md @@ -4,6 +4,50 @@ A running summary documenting some experiments and findings. Started ~Jan 7 2026 --- +## 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: From e3f58b838e98a5ea013a3c1773fde9d4a3c5d090 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 20:59:42 +0000 Subject: [PATCH 12/33] ranked version --- nanochat/gpt.py | 48 ++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 40 insertions(+), 8 deletions(-) diff --git a/nanochat/gpt.py b/nanochat/gpt.py index 86f440bf..ffb7862a 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -68,7 +68,7 @@ class CausalSelfAttention(nn.Module): self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False) self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False) - def forward(self, x, cos_sin, window_size, kv_cache): + def forward(self, x, cos_sin, window_size, kv_cache, v0, v0_lambda): B, T, C = x.size() # Project the input to get queries, keys, and values @@ -77,6 +77,11 @@ class CausalSelfAttention(nn.Module): 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 projected initial embedding for later layers + if v0 is not None: + v0_reshaped = v0.view(B, T, self.n_kv_head, self.head_dim) + v = v + v0_lambda * v0_reshaped + # 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) @@ -126,8 +131,8 @@ class Block(nn.Module): self.attn = CausalSelfAttention(config, layer_idx) self.mlp = MLP(config) - def forward(self, x, cos_sin, window_size, kv_cache): - x = x + self.attn(norm(x), cos_sin, window_size, kv_cache) + def forward(self, x, cos_sin, window_size, kv_cache, v0, v0_lambda): + x = x + self.attn(norm(x), cos_sin, window_size, kv_cache, v0, v0_lambda) x = x + self.mlp(norm(x)) return x @@ -160,6 +165,17 @@ class GPT(nn.Module): # 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() + # Value residual (ResFormer-style): low-rank factorized embedding for value residual + # Paper: "Value Residual Learning" (arXiv:2410.17897) shows this improves information flow + # We apply to last 1/4 of layers as the paper shows later layers benefit most + # Low-rank factorization: (vocab, r) @ (r, kv_dim) instead of full (vocab, kv_dim) + head_dim = config.n_embd // config.n_head + kv_dim = config.n_kv_head * head_dim + value_rank = 32 # low-rank bottleneck dimension + self.value_embed_A = nn.Embedding(padded_vocab_size, value_rank) # token -> low-rank + self.value_embed_B = nn.Linear(value_rank, kv_dim, bias=False) # low-rank -> kv_dim + self.v0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) # fake init, real init in init_weights() + self.value_residual_start = config.n_layer - config.n_layer // 4 # last 1/4 of layers # 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. @@ -204,15 +220,21 @@ class GPT(nn.Module): with torch.no_grad(): self.resid_lambdas.fill_(1.0) # 1.0 => typical residual connections at init self.x0_lambdas.fill_(0.0) # 0.0 => skip connection to input is disabled at init + self.v0_lambdas.fill_(0.0) # 0.0 => value residual is disabled at init + + # Value embedding low-rank factors (init like embeddings/projections) + torch.nn.init.normal_(self.value_embed_A.weight, mean=0.0, std=1.0) # like wte + torch.nn.init.uniform_(self.value_embed_B.weight, -s, s) # like c_v # 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 token embeddings to bf16: optimizer can tolerate it and it saves memory + # Cast embeddings to bf16: optimizer can tolerate it and it saves memory if self.transformer.wte.weight.device.type == "cuda": self.transformer.wte.to(dtype=torch.bfloat16) + self.value_embed_A.to(dtype=torch.bfloat16) def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None): # TODO: bump base theta more? e.g. 100K is more common more recently @@ -277,7 +299,8 @@ class GPT(nn.Module): """ nparams = sum(p.numel() for p in self.parameters()) # Exclude non-matmul params: embeddings and per-layer scalars - nparams_exclude = self.transformer.wte.weight.numel() + self.resid_lambdas.numel() + self.x0_lambdas.numel() + nparams_exclude = (self.transformer.wte.weight.numel() + self.value_embed_A.weight.numel() + + self.resid_lambdas.numel() + self.x0_lambdas.numel() + self.v0_lambdas.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 @@ -303,13 +326,16 @@ class GPT(nn.Module): def setup_optimizers(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0, adam_betas=(0.8, 0.95), scalar_lr=0.5): model_dim = self.config.n_embd ddp, rank, local_rank, world_size = get_dist_info() - # Separate out all parameters into 5 groups (matrix, embedding, lm_head, resid_lambdas, x0_lambdas) + # Separate out all parameters into groups (matrix, embedding, lm_head, value_embed, resid_lambdas, x0_lambdas, v0_lambdas) matrix_params = list(self.transformer.h.parameters()) embedding_params = list(self.transformer.wte.parameters()) lm_head_params = list(self.lm_head.parameters()) + value_embed_A_params = list(self.value_embed_A.parameters()) + value_embed_B_params = list(self.value_embed_B.parameters()) resid_params = [self.resid_lambdas] x0_params = [self.x0_lambdas] - assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(resid_params) + len(x0_params) + v0_params = [self.v0_lambdas] + assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(value_embed_A_params) + len(value_embed_B_params) + len(resid_params) + len(x0_params) + len(v0_params) # Create the AdamW optimizer for the embedding, lm_head, and per-layer scalars # Scale the LR for the AdamW parameters by ∝1/√dmodel (having tuned the LRs for 768 dim model) dmodel_lr_scale = (model_dim / 768) ** -0.5 @@ -317,8 +343,11 @@ class GPT(nn.Module): adam_groups = [ dict(params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale), dict(params=embedding_params, lr=embedding_lr * dmodel_lr_scale), + dict(params=value_embed_A_params, lr=embedding_lr * dmodel_lr_scale), # low-rank embedding + dict(params=value_embed_B_params, lr=embedding_lr * dmodel_lr_scale), # low-rank projection dict(params=resid_params, lr=scalar_lr * 0.01), # these are a lot more sensitive because they accumulate in the residual stream dict(params=x0_params, lr=scalar_lr), + dict(params=v0_params, lr=scalar_lr), ] adamw_kwargs = dict(betas=adam_betas, eps=1e-10, weight_decay=0.0) # NOTE: weight decay is hardcoded to 0.0 for AdamW, only used in Muon AdamWFactory = DistAdamW if ddp else partial(torch.optim.AdamW, fused=True) @@ -349,9 +378,12 @@ class GPT(nn.Module): x = self.transformer.wte(idx) x = norm(x) x0 = x # save initial normalized embedding for x0 residual + # Value residual (ResFormer): low-rank factorized embedding for later layers + v0 = self.value_embed_B(self.value_embed_A(idx)) # (B, T, kv_dim) for i, block in enumerate(self.transformer.h): x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0 - x = block(x, cos_sin, self.window_sizes[i], kv_cache) + v0_for_layer = v0 if i >= self.value_residual_start else None + x = block(x, cos_sin, self.window_sizes[i], kv_cache, v0_for_layer, self.v0_lambdas[i]) x = norm(x) # Forward the lm_head (compute logits) From 0b58d70e9975d42b4357dfb33f321f764759af9f Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 21:16:47 +0000 Subject: [PATCH 13/33] full ve version works very well --- nanochat/gpt.py | 28 +++++++++++----------------- 1 file changed, 11 insertions(+), 17 deletions(-) diff --git a/nanochat/gpt.py b/nanochat/gpt.py index ffb7862a..0356413d 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -165,15 +165,12 @@ class GPT(nn.Module): # 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() - # Value residual (ResFormer-style): low-rank factorized embedding for value residual + # Value residual (ResFormer-style): separate embedding for values, mixed into later layers # Paper: "Value Residual Learning" (arXiv:2410.17897) shows this improves information flow # We apply to last 1/4 of layers as the paper shows later layers benefit most - # Low-rank factorization: (vocab, r) @ (r, kv_dim) instead of full (vocab, kv_dim) head_dim = config.n_embd // config.n_head kv_dim = config.n_kv_head * head_dim - value_rank = 32 # low-rank bottleneck dimension - self.value_embed_A = nn.Embedding(padded_vocab_size, value_rank) # token -> low-rank - self.value_embed_B = nn.Linear(value_rank, kv_dim, bias=False) # low-rank -> kv_dim + self.value_embed = nn.Embedding(padded_vocab_size, kv_dim) self.v0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) # fake init, real init in init_weights() self.value_residual_start = config.n_layer - config.n_layer // 4 # last 1/4 of layers # To support meta device initialization, we init the rotary embeddings here, but it's just "fake" meta tensors only. @@ -222,9 +219,8 @@ class GPT(nn.Module): self.x0_lambdas.fill_(0.0) # 0.0 => skip connection to input is disabled at init self.v0_lambdas.fill_(0.0) # 0.0 => value residual is disabled at init - # Value embedding low-rank factors (init like embeddings/projections) - torch.nn.init.normal_(self.value_embed_A.weight, mean=0.0, std=1.0) # like wte - torch.nn.init.uniform_(self.value_embed_B.weight, -s, s) # like c_v + # Value embedding (init like c_v: uniform with same std) + torch.nn.init.uniform_(self.value_embed.weight, -s, s) # Rotary embeddings head_dim = self.config.n_embd // self.config.n_head @@ -234,7 +230,7 @@ class GPT(nn.Module): # Cast embeddings to bf16: optimizer can tolerate it and it saves memory if self.transformer.wte.weight.device.type == "cuda": self.transformer.wte.to(dtype=torch.bfloat16) - self.value_embed_A.to(dtype=torch.bfloat16) + self.value_embed.to(dtype=torch.bfloat16) def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None): # TODO: bump base theta more? e.g. 100K is more common more recently @@ -299,7 +295,7 @@ class GPT(nn.Module): """ nparams = sum(p.numel() for p in self.parameters()) # Exclude non-matmul params: embeddings and per-layer scalars - nparams_exclude = (self.transformer.wte.weight.numel() + self.value_embed_A.weight.numel() + + nparams_exclude = (self.transformer.wte.weight.numel() + self.value_embed.weight.numel() + self.resid_lambdas.numel() + self.x0_lambdas.numel() + self.v0_lambdas.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 @@ -330,12 +326,11 @@ class GPT(nn.Module): matrix_params = list(self.transformer.h.parameters()) embedding_params = list(self.transformer.wte.parameters()) lm_head_params = list(self.lm_head.parameters()) - value_embed_A_params = list(self.value_embed_A.parameters()) - value_embed_B_params = list(self.value_embed_B.parameters()) + value_embed_params = list(self.value_embed.parameters()) resid_params = [self.resid_lambdas] x0_params = [self.x0_lambdas] v0_params = [self.v0_lambdas] - assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(value_embed_A_params) + len(value_embed_B_params) + len(resid_params) + len(x0_params) + len(v0_params) + assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(value_embed_params) + len(resid_params) + len(x0_params) + len(v0_params) # Create the AdamW optimizer for the embedding, lm_head, and per-layer scalars # Scale the LR for the AdamW parameters by ∝1/√dmodel (having tuned the LRs for 768 dim model) dmodel_lr_scale = (model_dim / 768) ** -0.5 @@ -343,8 +338,7 @@ class GPT(nn.Module): adam_groups = [ dict(params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale), dict(params=embedding_params, lr=embedding_lr * dmodel_lr_scale), - dict(params=value_embed_A_params, lr=embedding_lr * dmodel_lr_scale), # low-rank embedding - dict(params=value_embed_B_params, lr=embedding_lr * dmodel_lr_scale), # low-rank projection + dict(params=value_embed_params, lr=embedding_lr * dmodel_lr_scale), # same LR as token embedding dict(params=resid_params, lr=scalar_lr * 0.01), # these are a lot more sensitive because they accumulate in the residual stream dict(params=x0_params, lr=scalar_lr), dict(params=v0_params, lr=scalar_lr), @@ -378,8 +372,8 @@ class GPT(nn.Module): x = self.transformer.wte(idx) x = norm(x) x0 = x # save initial normalized embedding for x0 residual - # Value residual (ResFormer): low-rank factorized embedding for later layers - v0 = self.value_embed_B(self.value_embed_A(idx)) # (B, T, kv_dim) + # Value residual (ResFormer): separate value embedding for later layers + v0 = self.value_embed(idx) # (B, T, kv_dim) for i, block in enumerate(self.transformer.h): x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0 v0_for_layer = v0 if i >= self.value_residual_start else None From 9a88194c3f684a3418c0c0f4069e6f3b3af10736 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 22:08:52 +0000 Subject: [PATCH 14/33] simply one VE per layer, works best --- nanochat/gpt.py | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/nanochat/gpt.py b/nanochat/gpt.py index 0356413d..ea7a4d86 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -165,14 +165,12 @@ class GPT(nn.Module): # 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() - # Value residual (ResFormer-style): separate embedding for values, mixed into later layers + # Value residual (ResFormer-style): every layer gets its own value embedding # Paper: "Value Residual Learning" (arXiv:2410.17897) shows this improves information flow - # We apply to last 1/4 of layers as the paper shows later layers benefit most head_dim = config.n_embd // config.n_head kv_dim = config.n_kv_head * head_dim - self.value_embed = nn.Embedding(padded_vocab_size, kv_dim) - self.v0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) # fake init, real init in init_weights() - self.value_residual_start = config.n_layer - config.n_layer // 4 # last 1/4 of layers + self.value_embeds = nn.ModuleList([nn.Embedding(padded_vocab_size, kv_dim) for _ in range(config.n_layer)]) + self.v0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) # fake init, real init in init_weights() # 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. @@ -219,8 +217,9 @@ class GPT(nn.Module): self.x0_lambdas.fill_(0.0) # 0.0 => skip connection to input is disabled at init self.v0_lambdas.fill_(0.0) # 0.0 => value residual is disabled at init - # Value embedding (init like c_v: uniform with same std) - torch.nn.init.uniform_(self.value_embed.weight, -s, s) + # Value embeddings (init like c_v: uniform with same std) + for ve in self.value_embeds: + torch.nn.init.uniform_(ve.weight, -s, s) # Rotary embeddings head_dim = self.config.n_embd // self.config.n_head @@ -230,7 +229,8 @@ class GPT(nn.Module): # Cast embeddings to bf16: optimizer can tolerate it and it saves memory if self.transformer.wte.weight.device.type == "cuda": self.transformer.wte.to(dtype=torch.bfloat16) - self.value_embed.to(dtype=torch.bfloat16) + for ve in self.value_embeds: + ve.to(dtype=torch.bfloat16) def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None): # TODO: bump base theta more? e.g. 100K is more common more recently @@ -295,7 +295,8 @@ class GPT(nn.Module): """ nparams = sum(p.numel() for p in self.parameters()) # Exclude non-matmul params: embeddings and per-layer scalars - nparams_exclude = (self.transformer.wte.weight.numel() + self.value_embed.weight.numel() + + value_embeds_numel = sum(ve.weight.numel() for ve in self.value_embeds) + nparams_exclude = (self.transformer.wte.weight.numel() + value_embeds_numel + self.resid_lambdas.numel() + self.x0_lambdas.numel() + self.v0_lambdas.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 @@ -322,15 +323,15 @@ class GPT(nn.Module): def setup_optimizers(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0, adam_betas=(0.8, 0.95), 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, embedding, lm_head, value_embed, resid_lambdas, x0_lambdas, v0_lambdas) + # Separate out all parameters into groups (matrix, embedding, lm_head, value_embeds, resid_lambdas, x0_lambdas, v0_lambdas) matrix_params = list(self.transformer.h.parameters()) embedding_params = list(self.transformer.wte.parameters()) lm_head_params = list(self.lm_head.parameters()) - value_embed_params = list(self.value_embed.parameters()) + value_embeds_params = list(self.value_embeds.parameters()) resid_params = [self.resid_lambdas] x0_params = [self.x0_lambdas] v0_params = [self.v0_lambdas] - assert len(list(self.parameters())) == len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(value_embed_params) + len(resid_params) + len(x0_params) + len(v0_params) + 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(v0_params) # Create the AdamW optimizer for the embedding, lm_head, and per-layer scalars # Scale the LR for the AdamW parameters by ∝1/√dmodel (having tuned the LRs for 768 dim model) dmodel_lr_scale = (model_dim / 768) ** -0.5 @@ -338,7 +339,7 @@ class GPT(nn.Module): adam_groups = [ dict(params=lm_head_params, lr=unembedding_lr * dmodel_lr_scale), dict(params=embedding_params, lr=embedding_lr * dmodel_lr_scale), - dict(params=value_embed_params, lr=embedding_lr * dmodel_lr_scale), # same LR as token embedding + dict(params=value_embeds_params, lr=embedding_lr * dmodel_lr_scale), # same LR as token embedding dict(params=resid_params, lr=scalar_lr * 0.01), # these are a lot more sensitive because they accumulate in the residual stream dict(params=x0_params, lr=scalar_lr), dict(params=v0_params, lr=scalar_lr), @@ -372,12 +373,11 @@ class GPT(nn.Module): x = self.transformer.wte(idx) x = norm(x) x0 = x # save initial normalized embedding for x0 residual - # Value residual (ResFormer): separate value embedding for later layers - v0 = self.value_embed(idx) # (B, T, kv_dim) + # Value residual (ResFormer): every layer gets its own value embedding + v0s = [ve(idx) for ve in self.value_embeds] # n_layer x (B, T, kv_dim) for i, block in enumerate(self.transformer.h): x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0 - v0_for_layer = v0 if i >= self.value_residual_start else None - x = block(x, cos_sin, self.window_sizes[i], kv_cache, v0_for_layer, self.v0_lambdas[i]) + x = block(x, cos_sin, self.window_sizes[i], kv_cache, v0s[i], self.v0_lambdas[i]) x = norm(x) # Forward the lm_head (compute logits) From e85db6b4a4351eb562bec220b3bbcaad28be6722 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Fri, 16 Jan 2026 23:52:12 +0000 Subject: [PATCH 15/33] alternating design --- nanochat/gpt.py | 60 +++++++++++++++++++++++++++---------------------- 1 file changed, 33 insertions(+), 27 deletions(-) diff --git a/nanochat/gpt.py b/nanochat/gpt.py index ea7a4d86..a077256e 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -45,6 +45,10 @@ def norm(x): return F.rms_norm(x, (x.size(-1),)) +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 @@ -67,8 +71,10 @@ class CausalSelfAttention(nn.Module): self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False) self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False) self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False) + self.ve_gate_channels = 32 + self.ve_gate = nn.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, cos_sin, window_size, kv_cache, v0, v0_lambda): + 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 @@ -77,10 +83,11 @@ class CausalSelfAttention(nn.Module): 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 projected initial embedding for later layers - if v0 is not None: - v0_reshaped = v0.view(B, T, self.n_kv_head, self.head_dim) - v = v + v0_lambda * v0_reshaped + # 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 = 2 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels])) # (B, T, n_kv_head), range (0, 2) + v = v + gate.unsqueeze(-1) * ve # Apply Rotary Embeddings to queries and keys to get relative positional encoding cos, sin = cos_sin @@ -131,8 +138,8 @@ class Block(nn.Module): self.attn = CausalSelfAttention(config, layer_idx) self.mlp = MLP(config) - def forward(self, x, cos_sin, window_size, kv_cache, v0, v0_lambda): - x = x + self.attn(norm(x), cos_sin, window_size, kv_cache, v0, v0_lambda) + 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 @@ -165,12 +172,10 @@ class GPT(nn.Module): # 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() - # Value residual (ResFormer-style): every layer gets its own value embedding - # Paper: "Value Residual Learning" (arXiv:2410.17897) shows this improves information flow + # 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.ModuleList([nn.Embedding(padded_vocab_size, kv_dim) for _ in range(config.n_layer)]) - self.v0_lambdas = nn.Parameter(torch.zeros(config.n_layer)) # fake init, real init in init_weights() + 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. @@ -181,6 +186,7 @@ class GPT(nn.Module): 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. @@ -212,15 +218,18 @@ class GPT(nn.Module): torch.nn.init.zeros_(block.mlp.c_proj.weight) # Per-layer scalars - with torch.no_grad(): - self.resid_lambdas.fill_(1.0) # 1.0 => typical residual connections at init - self.x0_lambdas.fill_(0.0) # 0.0 => skip connection to input is disabled at init - self.v0_lambdas.fill_(0.0) # 0.0 => value residual is disabled at init + self.resid_lambdas.fill_(1.0) # 1.0 => typical residual connections at init + self.x0_lambdas.fill_(0.0) # 0.0 => skip connection to input is disabled at init # Value embeddings (init like c_v: uniform with same std) - for ve in self.value_embeds: + for ve in self.value_embeds.values(): torch.nn.init.uniform_(ve.weight, -s, s) + # Gate weights init to zero so gates start at sigmoid(0) = 0.5, scaled by 2 -> 1.0 (neutral) + for block in self.transformer.h: + if block.attn.ve_gate is not None: + torch.nn.init.zeros_(block.attn.ve_gate.weight) + # Rotary embeddings head_dim = self.config.n_embd // self.config.n_head cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim) @@ -229,7 +238,7 @@ class GPT(nn.Module): # Cast embeddings to bf16: optimizer can tolerate it and it saves memory if self.transformer.wte.weight.device.type == "cuda": self.transformer.wte.to(dtype=torch.bfloat16) - for ve in self.value_embeds: + for ve in self.value_embeds.values(): ve.to(dtype=torch.bfloat16) def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None): @@ -295,9 +304,9 @@ class GPT(nn.Module): """ 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) + 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.v0_lambdas.numel()) + self.resid_lambdas.numel() + self.x0_lambdas.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 @@ -323,15 +332,14 @@ class GPT(nn.Module): def setup_optimizers(self, unembedding_lr=0.004, embedding_lr=0.2, matrix_lr=0.02, weight_decay=0.0, adam_betas=(0.8, 0.95), 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, embedding, lm_head, value_embeds, resid_lambdas, x0_lambdas, v0_lambdas) + # 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()) - value_embeds_params = list(self.value_embeds.parameters()) resid_params = [self.resid_lambdas] x0_params = [self.x0_lambdas] - v0_params = [self.v0_lambdas] - 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(v0_params) + 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) # Create the AdamW optimizer for the embedding, lm_head, and per-layer scalars # Scale the LR for the AdamW parameters by ∝1/√dmodel (having tuned the LRs for 768 dim model) dmodel_lr_scale = (model_dim / 768) ** -0.5 @@ -342,7 +350,6 @@ class GPT(nn.Module): dict(params=value_embeds_params, lr=embedding_lr * dmodel_lr_scale), # same LR as token embedding dict(params=resid_params, lr=scalar_lr * 0.01), # these are a lot more sensitive because they accumulate in the residual stream dict(params=x0_params, lr=scalar_lr), - dict(params=v0_params, lr=scalar_lr), ] adamw_kwargs = dict(betas=adam_betas, eps=1e-10, weight_decay=0.0) # NOTE: weight decay is hardcoded to 0.0 for AdamW, only used in Muon AdamWFactory = DistAdamW if ddp else partial(torch.optim.AdamW, fused=True) @@ -373,11 +380,10 @@ class GPT(nn.Module): x = self.transformer.wte(idx) x = norm(x) x0 = x # save initial normalized embedding for x0 residual - # Value residual (ResFormer): every layer gets its own value embedding - v0s = [ve(idx) for ve in self.value_embeds] # n_layer x (B, T, kv_dim) for i, block in enumerate(self.transformer.h): x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0 - x = block(x, cos_sin, self.window_sizes[i], kv_cache, v0s[i], self.v0_lambdas[i]) + ve = self.value_embeds[str(i)](idx) if str(i) in self.value_embeds else None + x = block(x, ve, cos_sin, self.window_sizes[i], kv_cache) x = norm(x) # Forward the lm_head (compute logits) From 3b95d4fd392fb4d593adb80530e80c8009d06f75 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 17 Jan 2026 00:23:30 +0000 Subject: [PATCH 16/33] allow label for scaling laws script --- scaling_laws.sh | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/scaling_laws.sh b/scaling_laws.sh index 321b286a..7c269c6a 100644 --- a/scaling_laws.sh +++ b/scaling_laws.sh @@ -1,5 +1,7 @@ #!/bin/bash +LABEL="jan16" + FLOPS_BUDGETS=( 1e18 3e18 @@ -7,14 +9,14 @@ FLOPS_BUDGETS=( ) DEPTHS=(8 10 12 14 16 18 20) NPROC_PER_NODE="${NPROC_PER_NODE:-8}" -WANDB_RUN="${WANDB_RUN:-scaling}" +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" +RESULTS_DIR="$NANOCHAT_BASE_DIR/scaling_laws_results_${LABEL}" mkdir -p "$RESULTS_DIR" RESULTS_FILE="$RESULTS_DIR/results.csv" From 1933e8504655b41626947ea5ef1addab6bfda236 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 17 Jan 2026 00:25:50 +0000 Subject: [PATCH 17/33] brief update to log --- dev/LOG.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/dev/LOG.md b/dev/LOG.md index ae518c83..c0ab680f 100644 --- a/dev/LOG.md +++ b/dev/LOG.md @@ -4,6 +4,18 @@ A running summary documenting some experiments and findings. Started ~Jan 7 2026 --- +## 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). From 6460dc6382a4f9dfd52d5f8db3b659b9674b47a9 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 17 Jan 2026 02:28:31 +0000 Subject: [PATCH 18/33] tweaks to readme a bit --- README.md | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index 9de2884c..9808c207 100644 --- a/README.md +++ b/README.md @@ -6,14 +6,15 @@ This repo is a full-stack implementation of an LLM like ChatGPT in a single, clean, minimal, hackable, dependency-lite codebase. nanochat is designed to run on a single 8XH100 node via scripts like [speedrun.sh](speedrun.sh), that run the entire pipeline start to end. This includes tokenization, pretraining, finetuning, evaluation, inference, and web serving over a simple UI so that you can talk to your own LLM just like ChatGPT. nanochat will become the capstone project of the course LLM101n being developed by Eureka Labs. +## Updates + +- (Jan 16 2026) The repo is in active development, I am currently fleshing out the pretraining stage. +- (Jan 7 2026) See new post: [nanochat Miniseries v1](https://github.com/karpathy/nanochat/discussions/420) and the associated script [miniseries.sh](miniseries.sh). + ## Talk to it To get a sense of the endpoint of this repo, you can currently find [nanochat d34](https://github.com/karpathy/nanochat/discussions/314) hosted on [nanochat.karpathy.ai](https://nanochat.karpathy.ai/). "d34" means that this model has 34 layers in the Transformer neural network. This model has 2.2 billion parameters, it was trained on 88 billion tokens by simply running the training script [run1000.sh](run1000.sh) with `--target_param_data_ratio=40` (2x longer than Chinchilla-optimal), and the total cost of training was ~$2,500 (about 100 hours training time on 8XH100 GPU node). While today this is enough to outperform GPT-2 of 2019, it falls dramatically short of modern Large Language Models like GPT-5. When talking to these micro models, you'll see that they make a lot of mistakes, they are a little bit naive and silly and they hallucinate a ton, a bit like children. It's kind of amusing. But what makes nanochat unique is that it is fully yours - fully configurable, tweakable, hackable, and trained by you from start to end. To train and talk to your own, we turn to... -## Updates - -- (Jan 7 2026) See new post: [nanochat Miniseries v1](https://github.com/karpathy/nanochat/discussions/420) and the associated script [miniseries.sh](miniseries.sh). - ## Quick start The fastest way to feel the magic is to run the speedrun script [speedrun.sh](speedrun.sh), which trains and inferences the $100 tier of nanochat. On an 8XH100 node at $24/hr, this gives a total run time of about 4 hours. Boot up a new 8XH100 GPU box from your favorite provider (e.g. I use and like [Lambda](https://lambda.ai/service/gpu-cloud)), and kick off the training script: @@ -99,7 +100,7 @@ And a bit more about computing environments that will run nanochat: ## Running on CPU / MPS -nanochat can be run on CPU or on MPS (if you're on Macbook), and will automatically try to detect what device is best to run on. You're not going to get too far without GPUs, but at least you'll be able to run the code paths and maybe train a tiny LLM with some patience. For an example of how to make all the run commands much smaller (feel free to tune!), you can refer to [dev/runcpu.sh](dev/runcpu.sh) file. You'll see that I'm essentially restricting all scripts to train smaller models, to run for shorter number of iterations, etc. This functionality is new, slightly gnarly (touched a lot of code), and was merged in this [CPU|MPS PR](https://github.com/karpathy/nanochat/pull/88) on Oct 21, 2025. +nanochat can be run on CPU or on MPS (if you're on Macbook) in principle, and will automatically try to detect what device is best to run on. The script [dev/runcpu.sh](dev/runcpu.sh) shows a very simple example that will exercise the code paths but basically produce garbage results. Unless you know what you're doing, I basically don't recommend using this script right now and hope to tune it a bit more in the future. ## Customization @@ -109,15 +110,9 @@ Additionally, to add new abilities to nanochat, see [Guide: counting r in strawb ## Questions -nanochat is designed to be short and sweet. One big advantage of this is that we can package up all of the files together and copy paste them to your favorite LLM to ask arbitrary questions. As an example, I like to package up the repo using the [files-to-prompt](https://github.com/simonw/files-to-prompt) utility like so: +I recommend using [DeepWiki](https://deepwiki.com/karpathy/nanochat) from Devin/Cognition to ask questions of this repo. In the URL of this repo, simply change github.com to deepwiki.com, and you're off. -```bash -files-to-prompt . -e py -e md -e html -e toml -e sh --cxml > packaged.txt -``` - -This includes all py, html, toml, sh files and chooses the cxml output format. Everything is written to the `packaged.txt` file, which atm measures ~330KB (i.e. well below ~100K tokens for a state of the art LLM), and ~8K lines of code in 45 files. - -Alternatively, I recommend using [DeepWiki](https://deepwiki.com/karpathy/nanochat) from Devin/Cognition to ask questions of this repo. In the URL of this repo, simply change github.com to deepwiki.com, and you're off. +You can also come to the [#nanochat Discord channel](https://discord.com/channels/1020383067459821711/1427295580895314031) to ask questions, or use the Discussions. ## Tests From e1dafc510f122d5c31c38a3c96e45e544f47930f Mon Sep 17 00:00:00 2001 From: Yamahammer <137644546+Yamahammer@users.noreply.github.com> Date: Fri, 16 Jan 2026 21:50:34 -0500 Subject: [PATCH 19/33] Reduce token waste in BOS bestfit by cropping shortest doc (#445) When no document fits the remaining row space, crop the shortest document in the buffer instead of the first. This minimizes discarded tokens. Co-authored-by: Claude Opus 4.5 --- nanochat/dataloader.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/nanochat/dataloader.py b/nanochat/dataloader.py index 562d517e..3e898935 100644 --- a/nanochat/dataloader.py +++ b/nanochat/dataloader.py @@ -178,8 +178,9 @@ def tokenizing_distributed_data_loader_with_state_bos_bestfit( doc = doc_buffer.pop(best_idx) row.extend(doc) else: - # No doc fits - crop first doc to fill remaining - doc = doc_buffer.pop(0) + # 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.extend(doc[:remaining]) rows.append(row[:row_capacity]) From f42ae9e901a34be3b06b6816d41871e54dedc986 Mon Sep 17 00:00:00 2001 From: Nitish Pandey <83660586+nitishpandey04@users.noreply.github.com> Date: Sat, 17 Jan 2026 08:26:43 +0530 Subject: [PATCH 20/33] fix condition to perform bpb evaluation (#324) Co-authored-by: svlandeg --- scripts/mid_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/mid_train.py b/scripts/mid_train.py index 01d9f7d4..c127c943 100644 --- a/scripts/mid_train.py +++ b/scripts/mid_train.py @@ -249,7 +249,7 @@ while True: last_step = bool(last_step_tensor.item()) # 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): + 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) From bbc4413c58bd24fed440030eee805121d5296340 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bar=C4=B1=C5=9F=20=C3=96zmen?= Date: Sat, 17 Jan 2026 05:59:12 +0300 Subject: [PATCH 21/33] Add high value engine tests for core invariants (33 LoC) (#396) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * test: add engine generation tests for expected invariants - test_seed_reproducibility - test_temperature_zero_determinism - test_max_tokens_respected - test_num_samples_count 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 * Fix temperature test * add test for seed variation in sampling Add test for seed variation in sampling with temperature > 0. * Rename test for clarity * Shorten assert msg --------- Co-authored-by: Claude Opus 4.5 Co-authored-by: Sofie Van Landeghem --- tests/test_engine.py | 69 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/tests/test_engine.py b/tests/test_engine.py index 9351e5a8..67b8a5c7 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -195,3 +195,72 @@ def test_multi_sample_first_token_diversity(): 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." From 77a46902e4557117d716bd9cd9604ec3913f25bd Mon Sep 17 00:00:00 2001 From: Yury Kirpichev Date: Fri, 16 Jan 2026 18:59:44 -0800 Subject: [PATCH 22/33] Fix WANDB_RUN parameter passing in runcpu.sh (#407) - Add --run=$WANDB_RUN to base_train, mid_train, and chat_sft calls - Ensures wandb logging works when WANDB_RUN environment variable is set - Matches the behavior in speedrun.sh Co-authored-by: svlandeg --- dev/runcpu.sh | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/dev/runcpu.sh b/dev/runcpu.sh index c0b32a54..6ed7a8af 100755 --- a/dev/runcpu.sh +++ b/dev/runcpu.sh @@ -41,7 +41,8 @@ python -m scripts.base_train \ --core-metric-every=50 \ --core-metric-max-per-task=12 \ --sample-every=50 \ - --num-iterations=50 + --num-iterations=50 \ + --run=$WANDB_RUN python -m scripts.base_loss --device-batch-size=1 --split-tokens=4096 python -m scripts.base_eval --max-per-task=16 @@ -52,7 +53,8 @@ python -m scripts.mid_train \ --eval-every=50 \ --eval-tokens=4096 \ --total-batch-size=1024 \ - --num-iterations=100 + --num-iterations=100 \ + --run=$WANDB_RUN # eval results will be terrible, this is just to execute the code paths. # note that we lower the execution memory limit to 1MB to avoid warnings on smaller systems python -m scripts.chat_eval --source=mid --max-new-tokens=128 --max-problems=20 @@ -63,7 +65,8 @@ python -m scripts.chat_sft \ --target-examples-per-step=4 \ --num-iterations=100 \ --eval-steps=4 \ - --eval-metrics-max-problems=16 + --eval-metrics-max-problems=16 \ + --run=$WANDB_RUN # Chat CLI # python -m scripts.chat_cli -p "Why is the sky blue?" From 2955650327fb71bc4a470d5e1093dd7c9cececfc Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 17 Jan 2026 03:16:12 +0000 Subject: [PATCH 23/33] add detection of device to report more correct mfu for bf16 --- nanochat/common.py | 49 +++++++++++++++++++++++++++++++++++++++++++ scripts/base_train.py | 11 +++++++--- 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/nanochat/common.py b/nanochat/common.py index 22559ce5..faf91440 100644 --- a/nanochat/common.py +++ b/nanochat/common.py @@ -200,3 +200,52 @@ class DummyWandb: pass def finish(self): pass + +# hardcoded BF16 peak flops for NVIDIA A100, H100, H200, B200 GPU and AMD MI250, MI300X, MI325X, MI355X and Intel PVC +# inspired by torchtitan: https://github.com/pytorch/torchtitan/blob/main/torchtitan/tools/utils.py +def get_peak_flops(device_name: str) -> float: + if "A100" in device_name: + # data from https://www.nvidia.com/en-us/data-center/a100/ + return 312e12 + elif "H100" in device_name: + # data from https://www.nvidia.com/en-us/data-center/h100/ + # NOTE: Specifications are one-half lower without sparsity. + if "NVL" in device_name: + return 835e12 + elif "PCIe" in device_name: + return 756e12 + else: # for H100 SXM and other variants + return 989e12 + elif "H200" in device_name: + # data from https://www.nvidia.com/en-us/data-center/h200/ + return 989e12 + elif "B200" in device_name: + # data from https://nvdam.widen.net/s/wwnsxrhm2w/blackwell-datasheet-3384703 + return 2.25e15 + elif "MI355X" in device_name: + # MI355X data from https://www.amd.com/en/products/accelerators/instinct/mi350/mi355x.html + return 2500e12 + elif "MI300X" in device_name or "MI325X" in device_name: + # MI300X data from https://www.amd.com/en/products/accelerators/instinct/mi300/mi300x.html + # MI325X data from https://www.amd.com/en/products/accelerators/instinct/mi300/mi325x.html + return 1300e12 + elif "MI250X" in device_name: + # data from https://www.amd.com/en/products/accelerators/instinct/mi200/mi250x.html (per GCD) + return 191.5e12 + elif "Data Center GPU Max 1550" in device_name: + # Also known as Ponte Vecchio (PVC). + # data from https://www.intel.com/content/www/us/en/docs/oneapi/optimization-guide-gpu/2025-0/intel-xe-gpu-architecture.html + # Dot Product Accumulate Systolic (DPAS): + # - Freq: 1300MHz + # - #ops: 512 + # Full EU mode (i.e. 512 max compute units): 340.8 TFLOPS (BF16) + # Standard EU mode (i.e. 448 max compute units): 298.2 TFLOPS (BF16) + max_comp_units = torch.xpu.get_device_properties("xpu").max_compute_units + return 512 * max_comp_units * 1300 * 10**6 + elif "l40s" in device_name: + # data from: "https://resources.nvidia.com/en-us-l40s/l40s-datasheet-28413" + return 362e12 + + else: # for other GPU types, assume A100 + logger.warning(f"Peak flops undefined for: {device_name}, fallback to A100") + return 312e12 diff --git a/scripts/base_train.py b/scripts/base_train.py index c61986e6..e051f99b 100644 --- a/scripts/base_train.py +++ b/scripts/base_train.py @@ -22,7 +22,7 @@ import torch from nanochat.gpt import GPT, GPTConfig 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 +from nanochat.common import compute_init, compute_cleanup, print0, DummyWandb, print_banner, get_base_dir, autodetect_device_type, get_peak_flops 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 @@ -82,6 +82,12 @@ master_process = ddp_rank == 0 # this process will do logging, checkpointing etc autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext() synchronize = torch.cuda.synchronize if device_type == "cuda" else lambda: None get_max_memory = torch.cuda.max_memory_allocated if device_type == "cuda" else lambda: 0 +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 @@ -395,8 +401,7 @@ while True: pct_done = 100 * step / num_iterations tok_per_sec = int(args.total_batch_size / dt) flops_per_sec = num_flops_per_token * args.total_batch_size / dt - promised_flops_per_sec_h100 = 989e12 * ddp_world_size # bfloat16 H100 SXM and without 2:4 sparsity - mfu = 100 * flops_per_sec / promised_flops_per_sec_h100 # in % + 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) From f5425245f99efd4145d2ac71a730af1e96777d6a Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 17 Jan 2026 03:22:20 +0000 Subject: [PATCH 24/33] more GPU types from PR 147 thanks @Qubitium --- nanochat/common.py | 109 ++++++++++++++++++++++++++++----------------- 1 file changed, 67 insertions(+), 42 deletions(-) diff --git a/nanochat/common.py b/nanochat/common.py index faf91440..44760f90 100644 --- a/nanochat/common.py +++ b/nanochat/common.py @@ -201,51 +201,76 @@ class DummyWandb: def finish(self): pass -# hardcoded BF16 peak flops for NVIDIA A100, H100, H200, B200 GPU and AMD MI250, MI300X, MI325X, MI355X and Intel PVC +# 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: - if "A100" in device_name: - # data from https://www.nvidia.com/en-us/data-center/a100/ - return 312e12 - elif "H100" in device_name: - # data from https://www.nvidia.com/en-us/data-center/h100/ - # NOTE: Specifications are one-half lower without sparsity. - if "NVL" in device_name: - return 835e12 - elif "PCIe" in device_name: - return 756e12 - else: # for H100 SXM and other variants - return 989e12 - elif "H200" in device_name: - # data from https://www.nvidia.com/en-us/data-center/h200/ - return 989e12 - elif "B200" in device_name: - # data from https://nvdam.widen.net/s/wwnsxrhm2w/blackwell-datasheet-3384703 + name = device_name.lower() + + # --- NVIDIA Blackwell --- + if "gb200" in name or "grace blackwell" in name: + return 2.5e15 + if "b200" in name: return 2.25e15 - elif "MI355X" in device_name: - # MI355X data from https://www.amd.com/en/products/accelerators/instinct/mi350/mi355x.html - return 2500e12 - elif "MI300X" in device_name or "MI325X" in device_name: - # MI300X data from https://www.amd.com/en/products/accelerators/instinct/mi300/mi300x.html - # MI325X data from https://www.amd.com/en/products/accelerators/instinct/mi300/mi325x.html - return 1300e12 - elif "MI250X" in device_name: - # data from https://www.amd.com/en/products/accelerators/instinct/mi200/mi250x.html (per GCD) - return 191.5e12 - elif "Data Center GPU Max 1550" in device_name: - # Also known as Ponte Vecchio (PVC). - # data from https://www.intel.com/content/www/us/en/docs/oneapi/optimization-guide-gpu/2025-0/intel-xe-gpu-architecture.html - # Dot Product Accumulate Systolic (DPAS): - # - Freq: 1300MHz - # - #ops: 512 - # Full EU mode (i.e. 512 max compute units): 340.8 TFLOPS (BF16) - # Standard EU mode (i.e. 448 max compute units): 298.2 TFLOPS (BF16) + if "b100" in name: + return 1.8e15 + + # --- NVIDIA Hopper (H100/H200/H800) --- + if "h200" in name: + if "nvl" in name or "pcie" in name: + return 836e12 + return 989e12 # H200 SXM + if "h100" in name: + if "nvl" in name: + return 835e12 + if "pcie" in name: + return 756e12 + return 989e12 # H100 SXM + if "h800" in name: + if "nvl" in name: + return 989e12 + return 756e12 # H800 PCIe + + # --- NVIDIA Ampere data center --- + if "a100" in name or "a800" in name: + return 312e12 + if "a40" in name: + return 149.7e12 + if "a30" in name: + return 165e12 + + # --- NVIDIA Ada data center --- + if "l40s" in name or "l40-s" in name or "l40 s" in name: + return 362e12 + if "l4" in name: + return 121e12 + + # --- AMD CDNA accelerators --- + if "mi355" in name: + return 2.5e15 + if "mi325" in name or "mi300x" in name: + return 1.3074e15 + if "mi300a" in name: + return 980.6e12 + if "mi250x" in name: + return 383e12 + if "mi250" in name: + return 362.1e12 + + # --- Intel --- + 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 - elif "l40s" in device_name: - # data from: "https://resources.nvidia.com/en-us-l40s/l40s-datasheet-28413" - return 362e12 - else: # for other GPU types, assume A100 - logger.warning(f"Peak flops undefined for: {device_name}, fallback to A100") - return 312e12 + # --- Consumer RTX (for hobbyists) --- + if "5090" in name: + return 209.5e12 + if "4090" in name: + return 165.2e12 + if "3090" in name: + return 71e12 + + # 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') From f9a7e0f111f9955a640c69cd3dfe457813dc4601 Mon Sep 17 00:00:00 2001 From: karpathy Date: Sat, 17 Jan 2026 12:27:30 -0800 Subject: [PATCH 25/33] update the CPU/MPS script to give reasonable results. The model can at least answer that Paris is the capital of France and knows that the sky is blue, for about 40 minutes of training on my macbook. Also fixed a bug that existed due to KVCache bfloat16 dtype assumption --- dev/runcpu.sh | 83 ++++++++++++++++++++------------------------ nanochat/engine.py | 11 +++++- scripts/base_loss.py | 17 ++++++++- tests/test_engine.py | 5 +-- 4 files changed, 67 insertions(+), 49 deletions(-) diff --git a/dev/runcpu.sh b/dev/runcpu.sh index 6ed7a8af..da8f6d19 100755 --- a/dev/runcpu.sh +++ b/dev/runcpu.sh @@ -1,12 +1,15 @@ #!/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 dev/cpu_demo_run.sh # NOTE: Training LLMs requires GPU compute and $$$. You will not get far on your Macbook. # Think of this run as educational/fun demo, not something you should expect to work well. -# This is also why I hide this script away in dev/ +# (This is why I hide this script away in dev/) +# You may also want to run this script manually and one by one, copy pasting commands into your terminal. # all the setup stuff export OMP_NUM_THREADS=1 @@ -20,58 +23,48 @@ if [ -z "$WANDB_RUN" ]; then WANDB_RUN=dummy fi -# wipe the report -python -m nanochat.report reset - -# train tokenizer on ~1B characters -python -m nanochat.dataset -n 4 -python -m scripts.tok_train --max-chars=1000000000 +# 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 very small 4 layer model on the CPU -# each optimization step processes a single sequence of 1024 tokens -# we only run 50 steps of optimization (bump this to get better results) +# 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=4 \ - --max-seq-len=1024 \ - --device-batch-size=1 \ - --total-batch-size=1024 \ - --eval-every=50 \ - --eval-tokens=4096 \ - --core-metric-every=50 \ - --core-metric-max-per-task=12 \ - --sample-every=50 \ - --num-iterations=50 \ + --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_loss --device-batch-size=1 --split-tokens=4096 +python -m scripts.base_loss --device-batch-size=1 --split-tokens=16384 python -m scripts.base_eval --max-per-task=16 -# midtraining +# midtraining (~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.mid_train \ - --max-seq-len=1024 \ - --device-batch-size=1 \ - --eval-every=50 \ - --eval-tokens=4096 \ - --total-batch-size=1024 \ - --num-iterations=100 \ - --run=$WANDB_RUN -# eval results will be terrible, this is just to execute the code paths. -# note that we lower the execution memory limit to 1MB to avoid warnings on smaller systems -python -m scripts.chat_eval --source=mid --max-new-tokens=128 --max-problems=20 - -# SFT -python -m scripts.chat_sft \ - --device-batch-size=1 \ - --target-examples-per-step=4 \ - --num-iterations=100 \ - --eval-steps=4 \ - --eval-metrics-max-problems=16 \ + --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 CLI -# python -m scripts.chat_cli -p "Why is the sky blue?" +# (it's ~ok to skip SFT) -# Chat Web -# python -m scripts.chat_web +# 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 -i mid -p "What is the capital of France?" -python -m nanochat.report generate +# Chat with the model over a pretty WebUI ChatGPT style +# python -m scripts.chat_web -i mid diff --git a/nanochat/engine.py b/nanochat/engine.py index 53fdec5b..7f05eb4e 100644 --- a/nanochat/engine.py +++ b/nanochat/engine.py @@ -90,7 +90,7 @@ class 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=torch.bfloat16): + 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 @@ -172,6 +172,13 @@ class Engine: """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) @@ -191,6 +198,7 @@ class Engine: batch_size=1, seq_len=len(tokens), device=device, + dtype=dtype, **kv_model_kwargs, ) ids = torch.tensor([tokens], dtype=torch.long, device=device) @@ -203,6 +211,7 @@ class Engine: batch_size=num_samples, seq_len=kv_length_hint, device=device, + dtype=dtype, **kv_model_kwargs, ) kv_cache_decode.prefill(kv_cache_prefill) diff --git a/scripts/base_loss.py b/scripts/base_loss.py index 6b44a30c..fb8cf596 100644 --- a/scripts/base_loss.py +++ b/scripts/base_loss.py @@ -104,7 +104,7 @@ for split_name in ["train", "val"]: bpb_results[split_name] = bpb print0(f"Model: {model_name}, {split_name} bpb: {bpb:.6f}") -# Master process also samples from the model (only for nanochat models) +# Master process also samples from the model for some basic knowledge-eliciting prompts (only for nanochat models) samples = [] if ddp_rank == 0 and args.hf_path is None: prompts = [ @@ -122,9 +122,23 @@ if ddp_rank == 0 and args.hf_path is None: with autocast_ctx: 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) +# Draw some unconditioned samples from the model (only for nanochat models) +unconditioned_samples = [] +if ddp_rank == 0 and args.hf_path is None: + engine = Engine(model, tokenizer) + tokens = tokenizer("", prepend="<|bos|>") + with autocast_ctx: + samples, _ = engine.generate_batch(tokens, num_samples=8, max_tokens=128, temperature=1.0) + for sample in samples: + sample_str = tokenizer.decode(sample) + print0("-" * 80) + print0(sample_str) + unconditioned_samples.append(sample_str) + # Log to report from nanochat.report import get_report get_report().log(section="Base model loss", data=[ @@ -134,6 +148,7 @@ get_report().log(section="Base model loss", data=[ "val bpb": bpb_results["val"], }, {f"sample {i}": sample for i, sample in enumerate(samples)}, + {f"unconditioned sample {i}": sample for i, sample in enumerate(unconditioned_samples)}, ]) # Cleanup diff --git a/tests/test_engine.py b/tests/test_engine.py index 67b8a5c7..01591111 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -96,6 +96,7 @@ def test_kv_cache_basic(): head_dim=head_dim, num_layers=num_layers, device="cpu", + dtype=torch.float32, ) # Check initial state @@ -130,7 +131,7 @@ def test_kv_cache_prefill(): # 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", + 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 @@ -140,7 +141,7 @@ def test_kv_cache_prefill(): # 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", + head_dim=head_dim, num_layers=num_layers, device="cpu", dtype=torch.float32, ) # Prefill From e7ed2082b836ac21e45020759e799c3bf1d511fe Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 17 Jan 2026 21:16:46 +0000 Subject: [PATCH 26/33] update the default GPTConfig kwargs otherwise they are confusing --- nanochat/gpt.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nanochat/gpt.py b/nanochat/gpt.py index a077256e..cb4bd05b 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -28,8 +28,8 @@ from nanochat.flash_attention import flash_attn @dataclass class GPTConfig: - sequence_len: int = 1024 - vocab_size: int = 50304 + 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) @@ -37,7 +37,7 @@ class GPTConfig: # Sliding window attention pattern string, tiled across layers. Final layer always L. # Characters: L=long (full context), S=short (half context) # Examples: "L"=all full context, "SL"=alternating, "SSL"=two short then one long - window_pattern: str = "L" + window_pattern: str = "SSSL" def norm(x): From 413e91aa0f5f3f841dbdc0009e64811cf75c5a9d Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sat, 17 Jan 2026 23:51:09 +0000 Subject: [PATCH 27/33] optimal ratio is now around 4 --- scripts/base_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/base_train.py b/scripts/base_train.py index c61986e6..bb8d8a68 100644 --- a/scripts/base_train.py +++ b/scripts/base_train.py @@ -47,7 +47,7 @@ parser.add_argument("--window-pattern", type=str, default="SSSL", help="sliding # 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=int, default=8, help="calculate num_iterations to maintain data:param ratio (Chinchilla=20, -1 = disable)") +parser.add_argument("--target-param-data-ratio", type=int, default=4, 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") parser.add_argument("--total-batch-size", type=int, default=524288, help="total batch size in tokens") From cf5c9e5b8eb2e06c7c2c1c4a280ed95a7f4aa68d Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 18 Jan 2026 00:07:08 +0000 Subject: [PATCH 28/33] resolve a crash for odd depths because FA3 needs head_dim % 8 == 0 --- scripts/base_train.py | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/scripts/base_train.py b/scripts/base_train.py index bb8d8a68..bcbd4841 100644 --- a/scripts/base_train.py +++ b/scripts/base_train.py @@ -106,21 +106,19 @@ vocab_size = tokenizer.get_vocab_size() print0(f"Vocab size: {vocab_size:,}") # Model kwargs are derived from the desired depth of the model +# We nudge model_dim up to the nearest multiple of head_dim to ensure clean division +# (FA3 requires head_dim divisible by 8, and this guarantees head_dim == args.head_dim exactly) +# (For very small depths, this gives a slight "unfair" advantage to models with odd depths) num_layers = args.depth -model_dim = args.depth * args.aspect_ratio -def find_num_heads(model_dim, target_head_dim): - # Find num_heads that divides model_dim evenly, with head_dim closest to target. - ideal = max(1, round(model_dim / target_head_dim)) - for offset in range(model_dim): - for candidate in [ideal + offset, ideal - offset]: - if candidate > 0 and model_dim % candidate == 0: - return candidate - return 1 -num_heads = find_num_heads(model_dim, args.head_dim) +base_dim = args.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 num_kv_heads = num_heads # default is 1:1 GQA (Group Query Attention) ratio (i.e. GQA is disabled) +head_dim = model_dim // num_heads print0(f"num_layers: {num_layers}") -print0(f"model_dim: {model_dim}") +print0(f"model_dim: {model_dim} (base: {base_dim}, nudge: {model_dim - base_dim:+d})") print0(f"num_heads: {num_heads}") +print0(f"head_dim: {head_dim}") print0(f"num_kv_heads: {num_kv_heads}") # Optimizer / data / training length related hyperparameters From babde18ce1cb59cb3d36f8874d1248983c7ba9c3 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 18 Jan 2026 03:00:38 +0000 Subject: [PATCH 29/33] small tweaks --- miniseries.sh | 1 - scaling_laws.sh | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/miniseries.sh b/miniseries.sh index 9a4512b6..c42544e3 100644 --- a/miniseries.sh +++ b/miniseries.sh @@ -61,7 +61,6 @@ for d in "${DEPTHS[@]}"; do # No --target-flops, let it use the default ratio from base_train torchrun --standalone --nproc_per_node=$NPROC_PER_NODE -m scripts.base_train -- \ --depth=$d \ - --target-param-data-ratio=8 \ --run="${WANDB_RUN}_d${d}" \ --model-tag="${TAG}" \ --core-metric-every=999999 \ diff --git a/scaling_laws.sh b/scaling_laws.sh index 7c269c6a..1f9dab87 100644 --- a/scaling_laws.sh +++ b/scaling_laws.sh @@ -7,7 +7,8 @@ FLOPS_BUDGETS=( 3e18 6e18 ) -DEPTHS=(8 10 12 14 16 18 20) +DEPTHS=(6 7 8 9 10 11 12 13 14) + 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) From d58fcd9d7331efba0224d59e026833738e7547a6 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 18 Jan 2026 03:01:13 +0000 Subject: [PATCH 30/33] log for jan 17 --- dev/LOG.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/dev/LOG.md b/dev/LOG.md index c0ab680f..8aeffdbf 100644 --- a/dev/LOG.md +++ b/dev/LOG.md @@ -4,6 +4,27 @@ A running summary documenting some experiments and findings. Started ~Jan 7 2026 --- +## 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. From 63bb5831e27ec4ad5f7493412cf16f3aa2a35877 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 18 Jan 2026 15:27:41 +0000 Subject: [PATCH 31/33] something i've wanted to do for a while - move all .sh runs to their own directory so they don't pollute root dir --- README.md | 27 ++++++++++++++----------- miniseries.sh => runs/miniseries.sh | 0 run1000.sh => runs/run1000.sh | 0 {dev => runs}/runcpu.sh | 0 scaling_laws.sh => runs/scaling_laws.sh | 0 speedrun.sh => runs/speedrun.sh | 0 6 files changed, 15 insertions(+), 12 deletions(-) rename miniseries.sh => runs/miniseries.sh (100%) rename run1000.sh => runs/run1000.sh (100%) rename {dev => runs}/runcpu.sh (100%) rename scaling_laws.sh => runs/scaling_laws.sh (100%) rename speedrun.sh => runs/speedrun.sh (100%) diff --git a/README.md b/README.md index 9808c207..fb8747fd 100644 --- a/README.md +++ b/README.md @@ -4,29 +4,29 @@ > The best ChatGPT that $100 can buy. -This repo is a full-stack implementation of an LLM like ChatGPT in a single, clean, minimal, hackable, dependency-lite codebase. nanochat is designed to run on a single 8XH100 node via scripts like [speedrun.sh](speedrun.sh), that run the entire pipeline start to end. This includes tokenization, pretraining, finetuning, evaluation, inference, and web serving over a simple UI so that you can talk to your own LLM just like ChatGPT. nanochat will become the capstone project of the course LLM101n being developed by Eureka Labs. +This repo is a full-stack implementation of an LLM like ChatGPT in a single, clean, minimal, hackable, dependency-lite codebase. nanochat is designed to run on a single 8XH100 node via scripts like [speedrun.sh](runs/speedrun.sh), that run the entire pipeline start to end. This includes tokenization, pretraining, finetuning, evaluation, inference, and web serving over a simple UI so that you can talk to your own LLM just like ChatGPT. nanochat will become the capstone project of the course LLM101n being developed by Eureka Labs. ## Updates - (Jan 16 2026) The repo is in active development, I am currently fleshing out the pretraining stage. -- (Jan 7 2026) See new post: [nanochat Miniseries v1](https://github.com/karpathy/nanochat/discussions/420) and the associated script [miniseries.sh](miniseries.sh). +- (Jan 7 2026) See new post: [nanochat Miniseries v1](https://github.com/karpathy/nanochat/discussions/420) and the associated script [miniseries.sh](runs/miniseries.sh). ## Talk to it -To get a sense of the endpoint of this repo, you can currently find [nanochat d34](https://github.com/karpathy/nanochat/discussions/314) hosted on [nanochat.karpathy.ai](https://nanochat.karpathy.ai/). "d34" means that this model has 34 layers in the Transformer neural network. This model has 2.2 billion parameters, it was trained on 88 billion tokens by simply running the training script [run1000.sh](run1000.sh) with `--target_param_data_ratio=40` (2x longer than Chinchilla-optimal), and the total cost of training was ~$2,500 (about 100 hours training time on 8XH100 GPU node). While today this is enough to outperform GPT-2 of 2019, it falls dramatically short of modern Large Language Models like GPT-5. When talking to these micro models, you'll see that they make a lot of mistakes, they are a little bit naive and silly and they hallucinate a ton, a bit like children. It's kind of amusing. But what makes nanochat unique is that it is fully yours - fully configurable, tweakable, hackable, and trained by you from start to end. To train and talk to your own, we turn to... +To get a sense of the endpoint of this repo, you can currently find [nanochat d34](https://github.com/karpathy/nanochat/discussions/314) hosted on [nanochat.karpathy.ai](https://nanochat.karpathy.ai/). "d34" means that this model has 34 layers in the Transformer neural network. This model has 2.2 billion parameters, it was trained on 88 billion tokens by simply running the training script [run1000.sh](runs/run1000.sh) with `--target_param_data_ratio=40` (2x longer than Chinchilla-optimal), and the total cost of training was ~$2,500 (about 100 hours training time on 8XH100 GPU node). While today this is enough to outperform GPT-2 of 2019, it falls dramatically short of modern Large Language Models like GPT-5. When talking to these micro models, you'll see that they make a lot of mistakes, they are a little bit naive and silly and they hallucinate a ton, a bit like children. It's kind of amusing. But what makes nanochat unique is that it is fully yours - fully configurable, tweakable, hackable, and trained by you from start to end. To train and talk to your own, we turn to... ## Quick start -The fastest way to feel the magic is to run the speedrun script [speedrun.sh](speedrun.sh), which trains and inferences the $100 tier of nanochat. On an 8XH100 node at $24/hr, this gives a total run time of about 4 hours. Boot up a new 8XH100 GPU box from your favorite provider (e.g. I use and like [Lambda](https://lambda.ai/service/gpu-cloud)), and kick off the training script: +The fastest way to feel the magic is to run the speedrun script [speedrun.sh](runs/speedrun.sh), which trains and inferences the $100 tier of nanochat. On an 8XH100 node at $24/hr, this gives a total run time of about 4 hours. Boot up a new 8XH100 GPU box from your favorite provider (e.g. I use and like [Lambda](https://lambda.ai/service/gpu-cloud)), and kick off the training script: ```bash -bash speedrun.sh +bash runs/speedrun.sh ``` Alternatively, since the script runs for 4 hours, I like to launch it like this inside a new screen session `speedrun` (and also log output to `speedrun.log`): ```bash -screen -L -Logfile speedrun.log -S speedrun bash speedrun.sh +screen -L -Logfile speedrun.log -S speedrun bash runs/speedrun.sh ``` See the [screen cheatsheet](https://gist.github.com/jctosta/af918e1618682638aa82) if you are less familiar. You can watch it go inside the screen session, or detach with `Ctrl-a d` and `tail speedrun.log` to view progress. Now wait 4 hours. Once it's done, you can talk to your LLM via the ChatGPT-like web UI. Make sure again that your local uv virtual environment is active (run `source .venv/bin/activate`), and serve it: @@ -73,7 +73,7 @@ Total wall clock time: 3h51m Unsurprisingly, $100 is not enough to train a highly performant ChatGPT clone. In fact, LLMs are famous for their multi-million dollar capex. For our purposes, I think there are two more scales of interest. First is the ~$300 tier d26 model (i.e. depth=26) that trains in ~12 hours, which slightly outperforms GPT-2 CORE score. Second is the $1000 tier (~41.6 hours), just because it's a nice round number. But both of these are not yet fully supported and therefore not attached here in the master branch yet. -That said, to give a sense, the example changes needed for the [speedrun.sh](speedrun.sh) file to train a GPT-2 grade model d26 only involve three changes: +That said, to give a sense, the example changes needed for the [speedrun.sh](runs/speedrun.sh) file to train a GPT-2 grade model d26 only involve three changes: ```bash ... @@ -100,7 +100,7 @@ And a bit more about computing environments that will run nanochat: ## Running on CPU / MPS -nanochat can be run on CPU or on MPS (if you're on Macbook) in principle, and will automatically try to detect what device is best to run on. The script [dev/runcpu.sh](dev/runcpu.sh) shows a very simple example that will exercise the code paths but basically produce garbage results. Unless you know what you're doing, I basically don't recommend using this script right now and hope to tune it a bit more in the future. +nanochat can be run on CPU or on MPS (if you're on Macbook) in principle, and will automatically try to detect what device is best to run on. The script [runcpu.sh](runs/runcpu.sh) shows a very simple example that will exercise the code paths but basically produce garbage results. Unless you know what you're doing, I basically don't recommend using this script right now and hope to tune it a bit more in the future. ## Customization @@ -132,8 +132,7 @@ python -m pytest tests/test_engine.py -v -s │ ├── gen_synthetic_data.py # Example synthetic data for identity │ ├── generate_logo.html │ ├── nanochat.png -│ ├── repackage_data_reference.py # Pretraining data shard generation -│ └── runcpu.sh # Small example of how to run on CPU/MPS +│ └── repackage_data_reference.py # Pretraining data shard generation ├── nanochat │ ├── __init__.py # empty │ ├── adamw.py # Distributed AdamW optimizer @@ -152,7 +151,12 @@ python -m pytest tests/test_engine.py -v -s │ ├── tokenizer.py # BPE Tokenizer wrapper in style of GPT-4 │ └── ui.html # HTML/CSS/JS for nanochat frontend ├── pyproject.toml -├── run1000.sh # Train the ~$800 nanochat d32 +├── runs +│ ├── miniseries.sh # Miniseries training script +│ ├── run1000.sh # Train the ~$800 nanochat d32 +│ ├── 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: calculate CORE score │ ├── base_loss.py # Base model: calculate bits per byte, sample @@ -165,7 +169,6 @@ python -m pytest tests/test_engine.py -v -s │ ├── mid_train.py # Chat model: midtraining │ ├── tok_eval.py # Tokenizer: evaluate compression rate │ └── tok_train.py # Tokenizer: train it -├── speedrun.sh # Train the ~$100 nanochat d20 ├── tasks │ ├── arc.py # Multiple choice science questions │ ├── common.py # TaskMixture | TaskSequence diff --git a/miniseries.sh b/runs/miniseries.sh similarity index 100% rename from miniseries.sh rename to runs/miniseries.sh diff --git a/run1000.sh b/runs/run1000.sh similarity index 100% rename from run1000.sh rename to runs/run1000.sh diff --git a/dev/runcpu.sh b/runs/runcpu.sh similarity index 100% rename from dev/runcpu.sh rename to runs/runcpu.sh diff --git a/scaling_laws.sh b/runs/scaling_laws.sh similarity index 100% rename from scaling_laws.sh rename to runs/scaling_laws.sh diff --git a/speedrun.sh b/runs/speedrun.sh similarity index 100% rename from speedrun.sh rename to runs/speedrun.sh From 6a477eedbdc8d2c66da84c2fbfcf907ee7e1ba60 Mon Sep 17 00:00:00 2001 From: xiayan0118 <49345397+xiayan0118@users.noreply.github.com> Date: Mon, 19 Jan 2026 17:19:51 -0800 Subject: [PATCH 32/33] fix: pass device_type to compute_init in engine.__main__ (#451) When running engine.py directly on non-GPU devices (CPU, MPS), compute_init() needs the device_type parameter to initialize correctly. This fixes failures on machines without CUDA support. --- nanochat/engine.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nanochat/engine.py b/nanochat/engine.py index 7f05eb4e..a1ba24c8 100644 --- a/nanochat/engine.py +++ b/nanochat/engine.py @@ -306,8 +306,8 @@ if __name__ == "__main__": """ import time # init compute - ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init() device_type = autodetect_device_type() + ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init(device_type) autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext() # load the model and tokenizer From 85b3e95e0966a9ef4d46c59c5598922a15affd51 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Sun, 25 Jan 2026 00:03:55 +0000 Subject: [PATCH 33/33] 320 experiments just to tune the adam beta1 of x0 a little bit up from 0.8 to 0.96 --- dev/LOG.md | 66 +++++++++++++++++++++++++++++++++++++++++++++++++ nanochat/gpt.py | 2 +- 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/dev/LOG.md b/dev/LOG.md index 8aeffdbf..068b35e9 100644 --- a/dev/LOG.md +++ b/dev/LOG.md @@ -4,6 +4,72 @@ A running summary documenting some experiments and findings. Started ~Jan 7 2026 --- +## 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: diff --git a/nanochat/gpt.py b/nanochat/gpt.py index cb4bd05b..f62d04be 100644 --- a/nanochat/gpt.py +++ b/nanochat/gpt.py @@ -349,7 +349,7 @@ class GPT(nn.Module): dict(params=embedding_params, lr=embedding_lr * dmodel_lr_scale), dict(params=value_embeds_params, lr=embedding_lr * dmodel_lr_scale), # same LR as token embedding dict(params=resid_params, lr=scalar_lr * 0.01), # these are a lot more sensitive because they accumulate in the residual stream - dict(params=x0_params, lr=scalar_lr), + dict(params=x0_params, lr=scalar_lr, betas=(0.96, 0.95)), # higher beta1 for x0 scalars ] adamw_kwargs = dict(betas=adam_betas, eps=1e-10, weight_decay=0.0) # NOTE: weight decay is hardcoded to 0.0 for AdamW, only used in Muon AdamWFactory = DistAdamW if ddp else partial(torch.optim.AdamW, fused=True)