Porting nanochat to a TPU: what carries over from PyTorch, and what breaks #1
tucan9389
announced in
Announcements
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
-
Karpathy's nanochat normally runs on an 8×H100 GPU node, and several ports of it to JAX already exist. Among them, my aim was to keep the config and architecture as close to nanochat as possible (parity) while catching up on both model quality and training performance — the quality (its CORE score) reproduced cleanly, and the performance came only partway. If nanochat is new to you: it's Karpathy's full-stack LLM project — tokenizer training, pretraining, SFT, and RL in a single repository — where about four hours and roughly $100 on one 8×H100 node gets you your own chatbot (nanochat is often called the "$100 speedrun"; those numbers are for d20). For reference, stopping at the GPT-2-grade base model (d24) takes about two hours and roughly $48 on the same node. This post is a record of that port: what carried over unchanged from PyTorch, and what broke on the TPU.
1. Reproduction results
nanochat-jax currently provides a speedrun.sh script that covers base model training and SFT (upstream nanochat goes all the way to RL; this reproduction stops at SFT). The claim above — that the quality reproduced — is based on the CORE score. CORE is the average of accuracies over 22 evaluation tasks, each rescaled so that random guessing scores 0 and a perfect score is 1. Within the same evaluation harness, it lets you compare models against each other, which is why nanochat uses it to judge what counts as "GPT-2 grade". What we reproduced is recipe 4 from the official nanochat LEADERBOARD (R4 from here on; it's also called d24, since depth 24 ≈ 1.4B parameters). The band (0.2512–0.2677) is the score distribution from Karpathy running the same R4 recipe 7 times; this run's 0.2695 lands just above it.
On performance, there's still a gap. MFU is about 24% (d24) — half of Karpathy's measured H100 numbers (47–48% at d20).
Below, we run each script of speedrun.sh on a TPU v6e-8 and check whether the quality Karpathy reported actually comes out.
CORE scores and the band are from the nanochat LEADERBOARD. w = wall clock including checkpointing and compilation; od = on-demand list price.1
If you spot anything wrong or unclear in this post, or have a question, please let me know at tucan.dev@gmail.com or in the comments — feedback is always welcome 🙂
2. TPU basics
This run used 8 v6e chips (a single-host slice). Per-chip specs across generations:
Sources: Cloud TPU System Architecture, v6e, tpu7x (Ironwood). bf16, Int8, and HBM numbers are from each generation's official spec sheet. Low-precision compute is roughly 2× bf16; v5p and v6e hardware-accelerate only up to Int8, and native FP8 support starts with Ironwood.
The v6e's MXU grew to 256×256, from the 128×128 of every generation up to v5p — if a tensor dimension isn't a multiple of 256, XLA pads it with zeros and part of the unit is wasted (more on this in insight 5 of section 4). Meanwhile, HBM is 32GB per chip — a third of v5p's 95GB — while compute is 2×: a compute-heavy, memory-lean design.
Our v6e-8 slice adds up to 7.34 bf16 PFLOPS and costs ~$4–5/hour on spot (us-central1). Spot is heavily discounted against on-demand, in exchange for GCP being able to reclaim (preempt) it at any time. This run cost $60.8 on spot over 12.19 hours total (~$263 at on-demand rates), with one preemption and recovery along the way. You're billed for as long as the node exists (training or not) — forget to delete it, and at $5/h for an 8-chip slice, a day is ≈ $120.
Figure 1. The JAX AI Stack — hardware (CPU/GPU/TPU) → the XLA compiler → the JAX core → the library layer on top (Flax and others). (Source: jaxstack.ai, © JAX team)
Of this stack, nanochat-jax uses three software layers — XLA, JAX, and Flax — plus Pallas, which isn't in the figure.
torch.nnjit·grad·vmap)autograd+torch.compiletorchdependency)torch.optimtorchdependency)torch.utils.data.pt(for nanochat compatibility).pt3. Speedrun verification — from the tokenizer to the report card
The pipeline runs tokenizer (5.3m) → base (6.02h + eval 44.5m) → SFT (68.7m + eval ~3.5h), and we run the scripts that speedrun.sh executes, one stage at a time (to run everything in one shot, see Appendix C). For a deep understanding of each stage, see Karpathy's walkthrough post (his is d20, ours is d24); this post focuses on quickly checking measured values against the reference numbers. The midtraining from the walkthrough era (an intermediate stage that pre-taught conversation format and tool use) no longer exists as a separate stage upstream (its data was folded into the SFT mixture), so we also go straight from base to SFT. Setup starts in Appendix A.
Figure 2. The nanochat speedrun pipeline — this reproduction runs tokenizer, base, and SFT (the solid boxes) and skips RL (dashed).
Step 1. Tokenizer (5.3m)
python -m nanochat_jax.dataset -n 170 # ClimbMix train 170 + val 1 shards python -m scripts.tok_train --max-chars 100000000 --doc-cap 10000 --vocab-size 32768 python -m scripts.tok_evalWe train a vocab-327682 tokenizer on 100M characters of ClimbMix, then compare against GPT-2 how many tokens the same text takes. On the training data (the
trainrow), ours uses 1.5% fewer tokens than the GPT-2 tokenizer. Fewer tokens means more text learned for the same budget — a compression advantage — and Karpathy's tokenizer shows the same pattern, a signal that the reproduction is on track (for Korean, the training data contains no Hangul, so ours spends about 2× the tokens of GPT-4; the full per-domain table is in Appendix D). The number to check:train+1.5%.Step 2. Base model (6.02h + 44.5m)
# --recipe through --use-real-data: Karpathy's R4 config / --attn-impl through --splash-*: v6e-8 TPU-specific python -m scripts.base_train \ --recipe=324e69c --depth=24 --seq-len=2048 --vocab-size=32768 \ --target-param-data-ratio=9.5 --total-batch-size=1048576 \ --device-batch-size=2 --grad-accum-steps=32 --grad-accum-impl=fused \ --warmup-steps=0 --warmdown-ratio=0.5 --final-lr-frac=0.0 \ --weight-decay=0.2 --matrix-lr=0.02 --embedding-lr=0.3 \ --unembedding-lr=0.004 --scalar-lr=0.5 \ --bf16 --cast-embeddings-bf16 --use-real-data \ --attn-impl=splash --splash-block-q=512 --splash-block-kv=512 --splash-block-kv-compute=256 \ --matmul-precision=default --lm-head-precision=highest --ve-grad-impl=onehot \ --checkpoint-every=200 --keep-last-checkpoints=2 \ --model-tag=d24_speedrun_r4 --no-final-evalThe top half of these arguments is the R4 recipe as-is (730M scaling parameters × 9.5 ≈ 6.9B tokens)3; the only ones you may need to touch are the
--splash-*kernel block sizes.4Against Karpathy R4's 2-hour base train loop (8×H100), ours takes 5.29h (6.02h wall clock including checkpointing and compilation), at about 24% MFU5.6 (The two lines above are an excerpt; I've uploaded the full training log of this run — the header also prints flops/token and peak TFLOPS.)
CORE (the 22-task metric from section 1) comes out to 0.2695 — above GPT-2's 0.2565, and in the same class as the R4 band (0.2512–0.2677). The val bpb of 0.7343 is a reference number only, since tokenizer differences mix into it (R4: 0.7185).
Figure 3. The dashed lines in the middle and right panels are the R4 baselines — by the end of training, both metrics reach the baselines, landing at R4 level (the curves are from an earlier reproduction run).
Step 3. SFT (68.7m + eval ~3.5h)
Train for 68.7 minutes on a mixture (~1.07M rows) of SmolTalk, MMLU, GSM8K, SpellingBee7, and 1,000 identity conversations8, and the base model that could only repeat the question back becomes a model that answers it. ChatCORE (the centered average over 6 chat tasks) is 0.3733 — the comparison line here is the base model itself. On the same 6 tasks, base scored essentially 0 on the generative ones, so most of the generative score is ability SFT added: SpellingBee 0.9961, GSM8K 0.1008 (math is RL's job). I don't compare 1:1 against the report card in Karpathy's walkthrough — that one is d20 and ran on the older midtraining pipeline.
Figure 4. Right panel: SFT (blue) is ahead of base (gray) on every task — you can see the generative tasks, which scored 0 on base, now score above zero.
Beyond the scores, checking with actual prompts:
Feed the same prompt to base and to SFT, and the difference is easy to see.
Report card
End to end: 12h14m wall clock (12.19h billed), $60.8 on spot (~$263 on-demand), one preemption with a 17-minute recovery (RL is out of scope, so its column stays empty — kept to match nanochat's format). The CORE reproduction is done. The next two sections cover the missteps on the way there, and what they cost.
4. Five things I learned from the TPU/JAX port
The most important thing the port taught me: if you can avoid it, don't. I decided to do it anyway, so here are the problems I ran into and some insights of my own, trimmed down to five. Each item is tagged with which side it hurt: quality (CORE) or performance (MFU). The first is the bug that fooled us the longest.
Insight 1. [Quality] Some bugs give wrong results without raising an error — data corruption from two lines of NumPy
When quality collapsed, the first thing I suspected was the chip, and the second was the bf16 embeddings. But what had quietly gone wrong was not the model — it was how the training loop fed the data: no error, no NaN, no shape mismatch, the kind of bug that shows up only as lower quality. Below is the code in base_train.py that collects the grad-accum microbatches.
Follow
train_loaderinto dataloader.py, and you findrow_buffer(line 112), which reuses a single buffer for speed and hands out views (zero-copy slices) of it. But the path above collects the microbatches in a list and callsnp.stacklater. Sincenp.asarraydoesn't copy something that is already an array,9 every collected microbatch ends up pointing at the last contents of the reused buffer — the batch for one step effectively becomes N copies of the last microbatch. That is data corruption: the token count stays the same, only the contents are wrong.The original PyTorch code draws one microbatch, immediately runs forward/backward to accumulate it into the gradients, then draws the next — so the data is consumed before the buffer is reused, and the trap doesn't exist there. We, on the other hand, collect all the microbatches and hand them to a single compiled step (
lax.scan) — and while they were being collected, the buffer got overwritten (why we needed this new path in the first place is explained by insight 2 below). The fix is one line: copy at capture time.The v6e runs from the period with this bug were stuck in a CORE band of 0.08–0.13; the run right after the fix scored 0.274 (that 0.274 is the published canonical run, separate from the speedrun's 0.2695 — both are R4-level). Also, the fix is already in main, so the commands you'd run from the speedrun section above are safe. (This data-feeding code actually had one more "sibling bug".10)
If a dataloader can hand out views of a reused buffer, then on any path that stacks them later, the safe first move is to check whether they get copied.
Insight 2. [Performance] It's better to merge the training step into a single jit
Port the PyTorch code as-is and you get the DON'T above. On GPU that structure is fine — asynchronous CUDA and
torch.compilehide the host-side dispatch cost. On TPU, an A/B of the two structures moves MFU from 16% to 22%.11The cause is the boundaries. XLA optimizes only within one jit program, so splitting the step into three programs leaves the device idle between programs (in our first implementation, idle time reached 26.7% of the step), and the compiler can't overlap operations across them either. The MaxText docs call wrapping the whole step in a single jit the "primary mechanism" behind their performance. Meanwhile, the first (manual) method in the official optax grad-accum example shows this structure with no warning, so it's easy to copy if you miss the scan alternative further down the page.
Rule: one compiled program per step. When performance is low, look at device idle time before you look at kernels (it shows in the XProf timeline) — if the idle time stands out, suspect the structure (the jit boundaries) first.
Insight 3. [Performance] The Splash attention kernel's defaults need tuning
python scripts/base_train.py ... \ --splash-block-q=512 --splash-block-kv=512 --splash-block-kv-compute=256 # DO — the values we use at seq 2048Splash is the TPU flash-attention kernel that ships inside JAX. Eight tile-size fields decide its performance — open the source and the defaults are all 128, with the authors' TODO still sitting right above them.
In an A/B that changed only the tiles, MFU went from 22.4% to 41.2%,11 and the loss matched to six decimal places — exactly what the docstring says: tiles barely touch the numerics and mostly move performance. Most of the runtime sat in the backward pass. Why flash-attention's backward costs more than its forward is covered in the FlashAttention-2 paper; the trade-offs of Splash tiles are in this write-up. MaxText also uses its own 512 instead of the JAX defaults.
The best values depend on the sequence length — in our seq-4096 experiments the optimum was 1024, and the published run (seq 2048) uses 512. The sweep from tpuchat, an independent port, converged to the same 1024 for its shapes — and 2048 tiles were actually slower. Bigger is not automatically better.
Since the loss isn't sensitive to the tiles (see the docstring), a tile sweep even before profiling is a reasonable first move.
Insight 4. [Quality] By default, TPUs compute fp32 matmuls in bf16
Our training command carries a flag called
--lm-head-precision highest. For a while I believed it was protecting quality; in reality, it does nothing. How that happened is tied to a TPU-specific behavior.Request an fp32 matmul on a TPU, and even though the dtype stays fp32, under the default settings the multiplication runs in bf16. This is not hidden behavior — it's a documented default — and it's a trap well known enough that a Google researcher wrote in an issue that his own team had been bitten by it for roughly the fourth time.
We actually hit this early in the port. One vocab projection amplified a 1.67e-6 input error by 197×, and a per-op
precision=higheston that operation stopped the divergence. Up to that point, it was a genuine fix. The mistake came after — we carried the flag over into bf16 training as well, but precision settings act only on fp32 operations (a maintainer's answer to exactly this question). Our layers cast the weights to bf16 before the multiply, so there is no fp32 operation left for the setting to act on.The quality gap happened to close in the same period, so the flag took the credit — but when I later retraced the commits, the most cleanly isolated cause was a different fix that landed alongside it: the missing grad-accum accumulation loop (the "sibling bug" from insight 1's footnote), not the precision flag. The flag stays in the command only to keep reproduction matched with the archived logs, and it's scheduled for removal in the next refactoring.
For work that does need fp32 precision (reference comparisons, fp32 evaluation, and so on), this default can be a problem. See: Cloud TPU bfloat16 docs, PyTorch/XLA precision tutorial.
Insight 5. [Performance] The best head_dim depends on the TPU generation — on v6e, multiples of 256 are recommended
Our model uses head_dim=128. Parity with nanochat was the goal, so we kept the reference value as-is. But how fast that value runs depends on the chip generation. The TPU matrix-multiply unit (MXU) is 256×256 on v6e and Ironwood, 128×128 on the generations before, and for a matmul to reach peak throughput, the dimensions need to be at least the MXU size. Accordingly, Google's model-design guide recommends sizing head_dim in multiples of 256 on Trillium (v6e). head_dim=128 was an exact fit up through v5p; on v6e, it no longer fills the MXU.
We measured MFU holding the parameter count equal (729.8M) and changing only the head decomposition, running each config for a short 12 steps (MFU stabilizes within a few steps). Read the difference between the two rows rather than the absolute values11:
That works out to about 17% faster at the same parameter count. Whether all of that difference comes from MXU alignment, though, we haven't isolated at the level of individual operations — the accurate statement stops at "a result consistent with feeding the MXU better". nanocode, a nanochat-based JAX/TPU project, chose 256 for the same reason (the snippet above).
For now we keep 128 for parity, but the architecture is open for later optimization work. One check remains first: we haven't A/B-tested whether quality (CORE) holds at the same parameter count — the speed is measured, the quality isn't.
The same config can align with the MXU differently depending on the chip generation.
Next, let me add up what all of this actually cost — missteps and trial and error included.
5. Why it cost about $2,500
Starting with almost no LLM pretraining experience, I spent three months and about $2,500 (TPU compute ~$2,350 + storage) getting to the $60.8 speedrun reproduction above. Even with coding agents doing a lot of the work, a bill this size honestly surprised me. The breakdown by category:
In short: the port itself (the bottom row) cost $40, and the rest went into reproducing the reference CORE with that port. Early on, I believed this recipe only ran on v5p. The v5p settings as-is hit OOM on v6e-8's HBM (32GB per chip vs v5p's 95GB); my current guess is that the microbatch size was the cause — at the time I changed the batch, the embeddings, and the training loop all at once during the migration, so I never isolated which change was decisive. After seeing nanocode's report of training a model of the same scale on a single v6e-8, I moved the batch and settings over to v6e, and one base training run dropped from ~$200 in the v5p days to ~$50, then down to ~$30 with the MFU optimizations (insights 2 and 3 of section 4). Today's $60.8 full run is that ~$30 base training plus tokenizer, SFT, and evaluation.
6. Wrapping up, and next steps
Thanks for following along this far. Quality has reached the same class on the CORE score; what remains is performance. MFU is up to about 24% (insights 2 and 3 of section 4 were the levers). Which of jit idle time, Splash tiles, head_dim alignment, or the attention HBM bottleneck dominates the remaining gap to Karpathy's 47–48% — that takes a roofline/XProf breakdown to answer, and that's the next piece of work. If you've dealt with this kind of problem before, I'd welcome your comments (MFU numbers from comparable cases are collected in the footnote5).
This is a starting point. Change the depth and build your own model, swap in a different tokenizer and dataset to teach it the knowledge you want, or pick up the optimization where we stopped. The $300 credit on a new GCP account covers three or four full runs on spot at about $61 each — but for a new account, requesting v6e quota comes first (Appendix B). If you'd rather play with the model before training anything, the base and SFT checkpoints from this run are up on Hugging Face.
7. Further reading
nanochat & LLM training basics
TPU & JAX
Neighboring projects
Repository links
Thanks to Jaewook Kang, Joosung Yoon, Taekmin Kim, Xing Liu, and one anonymous friend for reviewing this — and to my wife, for the budget and the time.
Appendix
Appendix A) Setup — from TPU allocation to environment
Allocate a v6e-8 TPU.
--spotrequests a spot instance.Once
READYappears, you can connect. It normally shows up within a few minutes, and you'll know about a failure within a minute.From
READYon, you're being billed ~$5/hour. Connect right away. (a) is an interactive shell; (b) and (c) are for one-shot commands and background runs.Once you're on the machine, clone the nanochat-jax repository and install the dependencies.
If you see eight TPU v6 lite devices, setup is done.
Appendix B) History and statistics of trying to allocate TPU spot instances
TPU spot attempt statistics by region, April–July 2026. This is an excerpt of the main chip–region combinations, so the rows add up to less than the full statistics (2,885 creation attempts, 1,632 failures).
Appendix C) Cheatsheet of GCP commands
In your local machine (Mac)
In your TPU instance (VM)
Appendix D) Full tok_eval compression tables
diff % = (baseline tokens − our tokens) ÷ baseline tokens. Positive means we fit the same bytes into fewer tokens (a compression advantage). ratio = bytes/token.
vs GPT-2
vs GPT-4
Footnotes
GPT-2's ~$43k is Karpathy's retrospective estimate (TPU v3 ×32 × 168h). The R4 row's ~2h and ~$48 are my estimate: a 2-hour base train loop times ~$24/hour for 8×H100. The ~$73 total is from Karpathy's 2026 tweet (GPT-2-grade d24, 8×H100, 3.04h) — that run is the default d24 (data-ratio 12), so it's the same d24 as our R4 (ratio 9.5) but a separate recipe. Our precise CORE value is 0.269486 (see the report card); the text uses 0.2695. ↩
The common belief that "rounding the vocab up to a hardware-friendly multiple makes it faster" comes from a GPU-era tweet by Karpathy (50257→50304, rounded up to a multiple of 64, ~25% faster via kernel occupancy). On TPU, XLA automatically pads matrix dimensions to multiples of the MXU tile (128 on v5p, 256 on v6e) (Cloud TPU performance guide), and 32768 is already such a multiple, so nothing is wasted on padding. On how vocab size relates to model scale, the vocab scaling law (Tao et al. 2024: larger models deserve larger vocabularies, but sublinearly) is a good read. ↩
The ~6.9B tokens come from the tokens:params ratio nanochat uses. It follows the compute-optimal ratio concept established by Chinchilla (the params–tokens combination that gives the lowest loss for a given compute budget), but the value itself is nanochat's own measurement, not Chinchilla's 20: an early scaling fit of ~8 (miniseries #420) was later re-measured to a compute-optimal ~10.5, which made the script default 12 (#481), and R4 uses 9.5, slightly undertrained relative to that. Karpathy's guess for why the value comes out below 20 is the influence of the Muon optimizer, among other things (#420). The ratio multiplies the non-embedding (scaling) parameters (d24 ≈ 730M): 730M × 9.5 ≈ 6.9B tokens. Of the ClimbMix shards, training uses about 150; 171 is the download count with some margin (170 train + 1 validation). ↩ ↩2
Block sizes are for performance and have almost no effect on numerics (Pallas
BlockSizesdocstring). The hard constraints mostly live in the Pallas splash kernel (JAX 0.10.0, the version we use), not in our code —block_kvmust be an integer multiple ofblock_kv_compute, andblock_kv_computemust be a multiple of 128 (=NUM_LANES) (kernel L967–970);block_qandblock_kvmust also be multiples of 128 (L667 and L673). Blocks must be ≤seq_len(L1446–1449) and must fit in on-chip VMEM (exceeding it gives a compile/allocation error — splash itself doesn't check this). The only constraint our wrapper enforces directly isblock_kv_compute % 128(attention.py#L82-L85). MaxText's default is also 512 (our starting value), and larger models raise onlyblock_q(gemma2-9b on v6e usessa_block_q=2048;block_kvstays at 512). Tune with XProf — the splash kernel tags its block sizes into the trace as XProf metadata. ↩In his homelab d24 reproduction (#710), matt-langston reports his DGX Spark at 29.6% MFU and puts the general H100 band at 50–65% (no depth given; the same thread also explains the difference between GPU utilization and MFU). Karpathy's own measured anchors are ~47–48% for the d20 speedrun and 47.40% for d34. nor, who ported modded-nanogpt to the same v6e-8, also landed at 15→23%, and the GPT-3 training MFU in the PaLM paper's table is 21.3%. So our 24% is not far outside this distribution. That said, these numbers differ in seq_len, model architecture, and how FLOPs are counted, so they can't be compared head-on — MFU was originally a metric for comparing system efficiency on the same model (PaLM §4.1). ↩ ↩2
For scaling and rooflines on TPU, Google DeepMind's scaling-book is a textbook-level resource. ↩
The general way to add narrow abilities like this via midtraining/SFT is laid out in Karpathy's "counting r in strawberry (and how to add abilities generally)" (nanochat #164). ↩
Karpathy, "Guide: infusing identity to your nanochat" (nanochat #139) — mixing 1,000 synthetic identity conversations into the midtraining/SFT mixture. ↩
np.asarraynot copying an input that is already an ndarray is NumPy's documented behavior — a slice is a view sharing the same buffer, and changes through a view show up in the original. NumPy inside data pipelines is a well-known source of silent bugs: Tanel Pärnamaa downloaded and audited 100K+ GitHub repositories that import PyTorch, and reported that among the repos combining a custom dataset, NumPy's random number generator, and multi-process loading, over 95% had a related bug (PyTorch's official tutorials, OpenAI, and NVIDIA code included; post). Not exactly the same bug as ours, but the same family. ↩An earlier version of the port had a "sibling bug", too — the grad-accum accumulation loop was missing entirely, so the model trained on only 1/N of the configured tokens while the LR schedule decayed as if full batches were being consumed (silent under-training — the early loss looks normal, and only the final quality suffers). Upstream nanochat derives
grad_accum_stepsfromtotal_batch_sizeand asserts the two divide evenly, which blocks this kind of mismatch at the source. Grad-accum is just an easy place for things to silently go wrong: HuggingFace and Unsloth also shipped grad-accum bugs (different mechanism — loss normalization) for a while in 2024 before fixing them (Unsloth and HuggingFace), and it's a known point of confusion in nanoGPT as well (#340). ↩The MFU comparisons in this section were measured on a comparison config aligned with nanocode (head_dim 256, seq_len 4096, value embeddings removed), to isolate just the structure and tile differences — which is why the absolute values differ from the published run (seq 2048, head 128, ~24% MFU). On this config, MFU climbed loop 16.2% → fused 22.4% → Splash tile tuning 41.2%, and the tiles (insight 3) were the bigger lever. The two structures in insight 2 can be reproduced with
--grad-accum-impl loop|fusedin our repo. The 52.5% nanocode reports on the same chip adds call-path and layout optimizations on top of this config, so it isn't directly comparable with our 24% (nanocode). The two MFU values in insight 5 (30.1% and 35.2%) were measured at seq 2048, the same as the published run. ↩ ↩2 ↩3 ↩4Beta Was this translation helpful? Give feedback.
All reactions