Porting nanochat to a TPU: what carries over from PyTorch, and what breaks · tucan9389/nanochat-jax · Discussion #1

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. Model Chip CORE (base) Train time Train cost Total cost GPT-2 (2019, 1.5B) TPU v3 x32 (est.) 0.2565 168 h ~$43,000 (est.) - nanochat R4 (d24) H100 x8 0.2571 ~2 h ~$48 (est.) - nanochat-jax (ours) TPU v6e-8 (spot) 0.2695 5.29 h (6.02 w) $30 ($130 od) $60.8 (~$263 od) 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: v5p v6e (Trillium) 7x (Ironwood) HBM capacity 95 GB 32 GB 192 GB HBM bandwidth 2,765 GB/s 1,638 GB/s 7,380 GB/s bf16 compute 459 TFLOPs 918 TFLOPs 2,307 TFLOPs Low-precision compute Int8 918 TOPs Int8

1,836 TOPs FP8 4,614 TFLOPs MXU (matrix-multiply unit) 128×128 256×256 256×256 Chips per host 4 8 4 1-host topology 2×2×1 2×4 2×2×1 Max pod 8,960 chips 256 chips 9,216 chips 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. Role nanochat-jax PyTorch equivalent Neural network modules Flax NNX torch.nn Computation + compilation JAX + XLA (jit·grad·vmap) eager + autograd + torch.compile Custom kernel Pallas (Splash Attention) Triton Optimizer JAX (own implementation, no torch dependency) torch.optim Data loading NumPy (no torch dependency) torch.utils.data Checkpoint PyTorch .pt (for nanochat compatibility) PyTorch .pt 3. 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_eval We 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 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)3; the only ones you may need to touch are the --splash-* kernel block sizes.4 [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% 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.) 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 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) # 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, 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: 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