Porting nanochat to a TPU: what carries over from PyTorch, and what breaks

A developer ported Karpathy's nanochat model to run on Google TPU v6e-8 using JAX, achieving comparable training quality but only half the hardware utilization (24% MFU vs. 47-48% on H100). The port required adapting PyTorch code for TPU-specific optimizations like Splash attention kernels and padding tensors to multiples of 256 for the MXU.

Porting nanochat to a TPU: what carries over from PyTorch, and what breaks 1 tucan9389 /tucan9389 announced in Announcements /tucan9389/nanochat-jax/discussions/categories/announcements - | 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 1. Reproduction resultsnanochat-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. 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. If you spot anything wrong or unclear in this post, or have a question, please let me know at 2. TPU basicsThis run used 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 Of this stack, nanochat-jax uses three software layers — XLA, JAX, and Flax — plus Pallas, which isn't in the figure. 3. Speedrun verification — from the tokenizer to the report cardThe 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 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 eval We train a vocab-32768 train row , 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-eval The top half of these arguments is the R4 recipe as-is 730M scaling parameters × 9.5 ≈ 6.9B tokens --splash- kernel block sizes. step 100 loss=4.517054 ... mfu=24.5% step 1000 loss=2.786205 ... mfu=24.7% Against 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% MFU The two lines above are an excerpt; I've uploaded the RESULTS DIR=$HOME/.cache/nanochat-jax/results/d24 speedrun r4 && mkdir -p $RESULTS DIR python -m scripts.base eval --source base --model-tag d24 speedrun r4 \ --eval bpb,core --eval-steps 20 --device-batch-size 1 --max-per-task -1 \ --out $RESULTS DIR/base eval.json --partial-out $RESULTS DIR/base eval core partial.json --bf16 CORE the 22-task metric from section 1 comes out to Step 3. SFT 68.7m + eval ~3.5h SFT training 68.7m python -m scripts.chat sft --model-tag d24 speedrun r4 --output-model-tag d24 speedrun r4 sft \ --num-iterations -1 --device-batch-size 1 --max-seq-len 2048 --total-batch-size 524288 \ --sft-scope full --eval-every 200 --eval-steps 4 \ --chatcore-every -1 --chatcore-max-cat -1 --chatcore-max-sample 24 \ --load-optimizer 1 --log-every 10 --bf16 scoring — the 3 multiple-choice tasks in one go, the 3 generative tasks one by one python -m scripts.chat eval -i sft -g d24 speedrun r4 sft --max-new-tokens 512 --batch-size 8 \ --dtype bfloat16 --out $RESULTS DIR/chat eval sft.json --task-name "ARC-Easy|ARC-Challenge|MMLU" for task in SpellingBee HumanEval GSM8K; do python -m scripts.chat eval -i sft -g d24 speedrun r4 sft --task-name "$task" \ --max-new-tokens 512 --dtype bfloat16 --num-samples 1 --temperature 0 --top-k 50 \ --jit-gen 1 --out $RESULTS DIR/chat eval $ echo "$task" | tr 'A-Z' 'a-z' .json done Train for 68.7 minutes on a mixture ~1.07M rows of SmolTalk, MMLU, GSM8K, SpellingBee , 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. Beyond the scores, checking with actual prompts: JAX PLATFORMS=cpu NANOCHAT JAX BASE DIR=<checkpoint dir PYTHONPATH=<nanochat-jax dir \ python scripts/chat cli.py -i sft -t 0.0 -p "How do I make coffee?" Feed the same prompt to base and to SFT, and the difference is easy to see. Prompt: How do I make coffee? BASE : How do I make coffee? How do I make coffee? ... repeats the question SFT : Making coffee is a simple process... you'll need a coffee maker, ... Prompt: Spell the word 'banana' letter by letter. BASE : The word 'banana' is made up of two letters. ... fails to spell SFT : banana:b,a,n,a,n,a Report card python -m nanochat jax.report generate | Metric | BASE | SFT | RL | |-----------------|----------|----------|----| | CORE metric | 0.269486 | - | - | | ARC-Challenge | - | 0.509386 | - | | ARC-Easy | - | 0.622896 | - | | GSM8K | - | 0.100834 | - | | HumanEval | - | 0.140244 | - | | MMLU | - | 0.369534 | - | | SpellingBee | - | 0.996094 | - | | train bpb | 0.731504 | - | - | | val bpb | 0.734295 | - | - | | ChatCORE metric | - | 0.373265 | - | End to end: 12h14m wall clock 12.19h billed , 4. Five things I learned from the TPU/JAX portThe 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. 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 idx chunks, target chunks = , for in range grad accum steps : inputs np, targets np = next train loader idx chunks.append np.asarray inputs np ← collects the views as-is the bug target chunks.append np.asarray targets np idx all = jnp.asarray np.stack idx chunks, axis=0 stacked later targets all = jnp.asarray np.stack target chunks, axis=0 Follow 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, idx chunks.append np.array inputs np, copy=True copy at capture time target chunks.append np.array targets np, copy=True 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". compute grads = jax.jit compute grads compile unit 1 accumulate = jax.jit accumulate compile unit 2 apply update = jax.jit apply update compile unit 3 acc = zero grads for mb in microbatches: DON'T — three programs chained by a Python loop acc = accumulate acc, compute grads state, mb state = apply update state, acc @jax.jit DO — the whole step is one program def train step state, microbatches : def micro step acc, mb : same logic as above — only the boundary moves inside return accumulate acc, compute grads state, mb , None acc = lax.scan micro step, zero grads, microbatches 0 return apply update state, acc Port the PyTorch code as-is and you get the DON'T above. On GPU that structure is fine — asynchronous CUDA and The 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 jax/experimental/pallas/ops/tpu/splash attention/splash attention kernel.py TODO apaszke,sharadmv : Select better parameters based on a heuristic. block q: int = 128 TRAP — all 8 tile fields default to 128 block kv: int = 128 block q dkv: int = 128 ...5 of them are backward-only 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 2048 Splash is the TPU flash-attention kernel that ships inside JAX. Eight tile-size fields decide its performance — open the In an A/B that changed only the tiles, MFU went from 22.4% to 41.2%, 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. Our training command carries a flag called Request an fp32 matmul on a TPU, and even though the dtype stays fp32, under the default settings the multiplication a = np.random.default rng 0 .standard normal 1024, 1024 .astype np.float32 ref = a.astype np.float64 @ a.T.astype np.float64 out = jax.jit jnp.dot a, a.T fp32 inputs, default precision assert np.abs np.asarray out, np.float64 - ref .max / np.abs ref .max < 1e-5 expected — passes on CPU; on TPU we expect it to fail, because of the default above We actually hit this early in the port. One vocab projection amplified a 1.67e-6 input error by 197×, and a per-op x = jnp.asarray a :8 , jnp.bfloat16 w = jnp.asarray a, jnp.bfloat16 weights get cast to bf16 before the multiply assert jnp.all jnp.dot x, w, precision="default" == jnp.dot x, w, precision="highest" TRAP — this passes. The two paths are bit-identical = evidence the setting is a no-op 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. nanocode/configs.py — d24: n embed 2048, n head 8 → head dim 256 "adjusted for head dim=256 to better feed the 256x256 tpuv6e systolic arrays" Our model uses head dim=128. 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 values 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". 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. Next, let me add up what all of this actually cost — missteps and trial and error included. 5. Why it cost about $2,500Starting 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: Quality regression after moving to v6e | ~$275 | The training loop rewritten for the v6e move the fused structure of section 4, insight 2 carried the NumPy bug from insight 1, so every step's batch collapsed into copies of the last microbatch and CORE dropped sharply; I misattributed it twice — to the chip, then to bf16 embeddings — and it took ten days to find | Five operational incidents | ~$585 | Lost checkpoints, overfitting from too little data, billing on a dead node, a zero-byte ckpt crash, and a multi-host launcher incident — five in total; prevention went into code atomic saves and operating rules Appendix C | Performance MFU optimization | ~$375 | The work behind the published run's ~24% MFU insights 2 and 3 of section 4 are its results . The big money went into two full trains launched for speed alone, without a quality gate | Publication rehearsal and regression hunting | ~$165 | Pre-release end-to-end runs fixed 3 problems that only show up in a full training run a disk blow-up, eval attention, SFT OOM | Finishing the final reproduction | ~$180 | The first v6e reproduction, ChatEval, and several final clean-clone runs — the last $60.8 speedrun reached CORE 0.2695, the level of nanochat's R4 | Port correctness checks and misc. | ~$40 | This went into checking that runs split across 1–16 chips produce the same results, plus evaluation; the PyTorch tensor comparisons finished on a laptop CPU, at no TPU cost | 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 6. Wrapping up, and next stepsThanks 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 footnote 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 7. Further reading Appendix Appendix A Setup — from TPU allocation to environmentAllocate a v6e-8 TPU. gcloud alpha compute tpus tpu-vm create <N \ --zone=<Z --accelerator-type=v6e-8 --version=v6e-ubuntu-2404 --spot Once From a interactive shell — for quick commands gcloud compute tpus tpu-vm ssh <N --zone=<Z b one-shot command — put multi-line scripts in the REMOTE variable via a quoted heredoc REMOTE=$ cat <<'EOF' <script to run on the VM EOF gcloud compute tpus tpu-vm ssh <N --zone=<Z --ssh-flag="-oConnectTimeout=30" --command="$REMOTE" c long-running command — detach from the session with setsid nohup and run in the background REMOTE=$ cat <<'EOF' setsid nohup bash -c '<script to run on the VM ' ~/logs/command0.log 2 &1 & echo LAUNCHED EOF gcloud compute tpus tpu-vm ssh <N --zone=<Z --ssh-flag="-oConnectTimeout=30" --command="$REMOTE" Once you're on the machine, clone the nanochat-jax repository and install the dependencies. sudo rm -f /tmp/libtpu lockfile required before the first jax process git clone https://github.com/tucan9389/nanochat-jax.git curl -LsSf https://astral.sh/uv/install.sh | sh uv venv --python 3.12 ~/venv && source ~/venv/bin/activate uv pip install "torch~=2.11.0" --index-url https://download.pytorch.org/whl/cpu cd ~/nanochat-jax && uv pip install -e ". tpu,dev " -f https://storage.googleapis.com/jax-releases/libtpu releases.html python -c "import jax; print jax.devices " If you see eight TPU v6 lite devices, setup is done. Appendix B History and statistics of trying to allocate TPU spot instancesTPU 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 allocate TPU — sweep zones until you land on one with capacity success rates and prices: Appendix B for z in us-central1-b us-central1-a us-central1-c us-east4-a asia-southeast1-b; do gcloud alpha compute tpus tpu-vm create <N \ --zone=$z --accelerator-type=v6e-8 --version=v6e-ubuntu-2404 --spot --quiet && break done check the status of machine gcloud compute tpus tpu-vm describe <N --zone=<Z --format='value state ' READY / PREEMPTED gcloud compute tpus tpu-vm list --zone=<Z run a command on the TPU machine gcloud compute tpus tpu-vm ssh <N --zone=<Z interactive shell gcloud compute tpus tpu-vm ssh <N --zone=<Z --command="<CMD " one-shot command for multi-line scripts, see heredoc patterns b c in Appendix A gcloud compute tpus tpu-vm scp <N :~/logs/31 base train.log ./logs/ --zone=<Z retrieve logs deallocate the TPU machine after deleting, sweep every zone you tried for zombie instances — create is async gcloud compute tpus tpu-vm delete <N --zone=<Z --quiet for z in us-central1-b us-central1-a us-central1-c us-east4-a asia-southeast1-b; do gcloud compute tpus tpu-vm list --zone=$z 2 &1 | grep -v "Listed 0" done anything else? gcloud storage buckets create gs://<BUCKET --location=<REGION same region as the TPU zone gcloud storage rsync -r gs://<BUCKET /base checkpoints /path/to/ssd/base checkpoints permanent mirror of the results gcloud auth login; gcloud config set project <PROJECT-ID first time only core command within speedrun.sh sudo rm -f /tmp/libtpu lockfile required before the first jax process bash runs/speedrun.sh the production pipeline in one shot bash runs/runcpu.sh d2/CPU wiring check only checks it runs to the end the stages speedrun.sh runs in order same as Steps 1-3 in the main text : python -m nanochat jax.dataset -n 8 B-1 for the tokenizer python -m nanochat jax.dataset -n 170 & B-2 for training background python -m scripts.tok train --max-chars 100000000 --doc-cap 10000 --vocab-size 32768 B-3 python -m scripts.tok eval B-4 python -m scripts.base train <Step 2 arguments from the main text B-5 python -m scripts.base eval --source base --model-tag d24 speedrun r4 --eval bpb,core ... B-6 python -m scripts.chat sft <Step 3 arguments from the main text B-7 python -m scripts.chat eval -i sft -g d24 speedrun r4 sft ... B-8 python -m nanochat jax.report generate B-9 monitoring base training — the GCS checks and preemption recovery below only work if you mirror checkpoints to <BUCKET periodic gcloud storage rsync recommended tail -5 ~/logs/31 base train.log | grep -E 'step|bpb' last step lines ps -ef | grep ' b ase train' process alive b avoids grep matching itself df -h / | tail -1 disk gcloud storage ls -l gs://<BUCKET /base checkpoints/d24 speedrun r4/ | tail -8 latest ckpt in GCS monitoring SFT training tail -5 ~/logs/60 sft chain.log | grep -E 'step|val bpb' ps -ef | grep ' c hat sft' anything else? — preemption recovery restore from GCS, then continue training: python -m scripts.base train ... --resume-from-step <N eval: python -m scripts.chat eval ... --start-index <number of completed problems Appendix D Full tok eval compression tablesdiff % = baseline tokens − our tokens ÷ baseline tokens. Positive means we fit the same bytes into fewer tokens a compression advantage . ratio = bytes/token. | Text | Bytes | GPT-2 tok | ratio | Ours tok | ratio | diff % | |---------|---------|-----------|-------|----------|-------|---------| | news | 1803 | 392 | 4.60 | 406 | 4.44 | -3.6% | | korean | 893 | 745 | 1.20 | 729 | 1.22 | +2.1% | | code | 1259 | 576 | 2.19 | 410 | 3.07 | +28.8% | | math | 690 | 348 | 1.98 | 338 | 2.04 | +2.9% | | science | 1098 | 253 | 4.34 | 238 | 4.61 | +5.9% | | train | 2948778 | 631304 | 4.67 | 621881 | 4.74 | +1.5% | | val | 3024593 | 653067 | 4.63 | 645862 | 4.68 | +1.1% | | Text | Bytes | GPT-4 tok | ratio | Ours tok | ratio | diff % | |---------|---------|-----------|-------|----------|-------|---------| | news | 1803 | 387 | 4.66 | 406 | 4.44 | -4.9% | | korean | 893 | 364 | 2.45 | 729 | 1.22 | -100.3% | | code | 1259 | 309 | 4.07 | 410 | 3.07 | -32.7% | | math | 690 | 295 | 2.34 | 338 | 2.04 | -14.6% | | science | 1098 | 245 | 4.48 | 238 | 4.61 | +2.9% | | train | 2948778 | 611619 | 4.82 | 621881 | 4.74 | -1.7% | | val | 3024593 | 631183 | 4.79 | 645862 | 4.68 | -2.3% | Footnotes | Beta Was this translation helpful? Give feedback.