tucan9389announced in
| 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)
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) #
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
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_)
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
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
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"))
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.
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
gcloud compute tpus tpu-vm ssh <N> --zone=<Z>
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"
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 #
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
gcloud compute tpus tpu-vm describe <N> --zone=<Z> --format='value(state)' # READY / PREEMPTED
gcloud compute tpus tpu-vm list --zone=<Z>
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
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
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
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)
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
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
tail -5 ~/logs/60_sft_chain.log | grep -E 'step|val_bpb'
ps -ef | grep '[c]hat_sft'
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.