nanochatをTPUに移植する：PyTorchから引き継がれるものと壊れるもの

tucan9389 / nanochat-jax Public Notifications You must be signed in to change notification settings Fork 0 Star 3 Porting nanochat to a TPU: what carries over from PyTorch, and what breaks #1 tucan9389 announced in Announcements Porting nanochat to a TPU: what carries over from PyTorch, and what breaks #1 tucan9389 Jul 15, 2026 · 0 comments Return to top Discussion options Uh oh! There was an error while loading. Please reload this page. {{title}} Something went wrong. Uh oh! There was an error while loading. Please reload this page. Quote reply tucan9389 Jul 15, 2026 Maintainer - 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. 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: Chip v5p v6e (Trillium) v7x (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