GPT-2 grade on a single RTX 5090 in 41 hours — recipe, Blackwell gotcha, and an open question about sliding windows

[carey-bunks]

GPT-2 grade on a single RTX 5090 in 41 hours — recipe, Blackwell gotcha, and an open question about sliding windows

Not a leaderboard entry — the leaderboard measures wall-clock on 8×H100 and this is one consumer card. Posting because the README points sub-80GB users at --device-batch-size and leaves them to it, and there were a few more things to get right than that.

Result

Hardware	1 × RTX 5090 (32 GB), sm_120
Model	d24, --target-param-data-ratio=8, 1.384B params
Wall clock	41.4 h pretraining (2481 min)
CORE (standalone base_eval)	0.2702
CORE (in-training, step 5568)	0.2803
val bpb / train bpb	0.720161 / 0.721011
Peak memory (torch-reported)	21,122 MiB
Step time	~27.1 s × 5,568 steps

Data: 170 ClimbMix shards downloaded, 131 consumed, epoch: 1 throughout. Tokenizer 4.74 chars/token. Cost: electricity (approx. £6.64 or $8.95).

The Blackwell gotcha (this one is mandatory)

On sm_120 the FA3 loader reports success and then dies at the first attention call:

✓ Using Flash Attention 3: efficient, new and awesome.
...
CUDA error (/build/source/flash-attn/flash_fwd_launch_template.h:192): no kernel image is available for execution on the device

_load_flash_attention_3() only special-cases major == 9 for Hopper; Blackwell falls through to kernels-community/flash-attn3, and has_kernel() returns True for a kernel with no sm_120 image. One-line fix in nanochat/flash_attention.py, right after the capability check:

        if major >= 10:
            return None   # Blackwell sm100/sm120: no FA3 kernels available

Verify:

python -c "from nanochat.flash_attention import USE_FA3, HAS_FA3; print(USE_FA3, HAS_FA3)"
# False False

(Filed separately as an issue.)

Full recipe

# environment — note NANOCHAT_BASE_DIR if you keep more than one install
export NANOCHAT_BASE_DIR="$HOME/.cache/nanochat-climbmix"
export PYTORCH_ALLOC_CONF=expandable_segments:True   # not PYTORCH_CUDA_ALLOC_CONF, deprecated
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS=1
uv venv && uv sync --extra gpu && source .venv/bin/activate
apply the flash_attention.py patch here
python -m nanochat.dataset -n 8 -w 2

python -m scripts.tok_train

python -m scripts.tok_eval

python -m nanochat.dataset -n 170 -w 2
torchrun --standalone --nproc_per_node=1 -m scripts.base_train -- 

--model-tag d24-climbmix --depth=24 --target-param-data-ratio=8 

--device-batch-size=2 --fp8 --window-pattern L
torchrun --standalone --nproc_per_node=1 -m scripts.base_eval -- 

--model-tag d24-climbmix --device-batch-size=16

Notes on each deviation

--device-batch-size=2 → 256 gradient accumulation steps. Peak was 21 GB torch-reported (~29.5 GB in nvidia-smi, the difference being allocator reservation). 4 does not fit for d24 on 32 GB. base_eval at 16 is fine — no optimizer state.

--window-pattern L — forced, not chosen. Sliding windows need FA3; without it _sdpa_attention has to build an explicit mask, which disables the fused kernel, and base_train.py warns about this directly. With L at Tq == Tk == 2048 the SDPA path returns via F.scaled_dot_product_attention(..., is_causal=True) with no mask, so you stay on torch's own fused flash backend, which does support sm_120.

--fp8 works on Blackwell — 145/158 linear layers converted, no issues across 41 hours.

Use torchrun even on one GPU. Bare python -m scripts.base_train hung before argparse and was unresponsive to SIGINT on this setup. The -- separator is required when going through torchrun.

Don't run runs/speedrun.sh — it hardcodes NANOCHAT_BASE_DIR, 8 GPUs, device-batch 16, and the default window pattern.

--save-every defaults to −1, i.e. checkpoints only at the end. I gambled and got away with it; on a 41-hour run you probably want --save-every=500. If you're on a consumer card, also check your case thermals rather than just the GPU — I had a prior kernel panic at 12 hours under sustained load at stock 575 W. Ran this one capped at 500 W, ~78 °C, with extra airflow, and step time held at 27.1 s for the full run with zero drift.

Reading the log: bf16_mfu shows ~17% but real utilisation is ~100%. Estimated FLOPs per token for d24 prints 9.06e8, which appears to count only the attention term — the matmul term is missing, true value ≈ 5.28e9. (38,800 tok/s × 5.28e9 = 2.05e14 FLOP/s against the card's 2.10e14 bf16 peak.) Trust the eta field, which is wall-clock derived.

An open question about sliding windows

My CORE came in above the ClimbMix leaderboard entry (0.2702 vs 0.2571) at essentially identical val_bpb (0.7188 in-training vs 0.71854). I want to be clear that this is not evidence for anything — I'm on master four months after that entry, at d24 rather than d26, ratio 8 rather than 8.25, device-batch 2 with 256 accumulation steps rather than 8×16, FP8 on a different hardware path, one seed, and I'm not certain whether leaderboard CORE figures come from the capped in-training eval or the full standalone one (that gap alone is 0.0101 for me).

But it does raise a question I can't answer from one run: what does the SSSL sliding-window pattern cost in quality? It's there to save compute, and on H100s with FA3 it clearly does. Whether the local-attention inductive bias is free, mildly positive, or mildly negative at fixed step count doesn't seem to have been isolated anywhere I could find.

That's a clean d12 A/B — L vs SSSL, everything else held — and cheap on 8×H100 at the ~5 min/run scale the README recommends for iteration. I can run it on the 5090 (slowly, since SSSL will be on the unfused SDPA path here), but someone with the reference hardware would get a much cleaner answer, and paired val_bpb + CORE would settle it. Happy to be told this has already been measured.

Other small things

• nvidia-smi overstates memory considerably vs the torch-reported peak; don't size batches from it
• Shard downloads fail intermittently; re-running the same command skips what's present and fetches the gaps. Drop -w to 2 if you see failures
• Piping to tee without PYTHONUNBUFFERED=1 gives a 0-byte log for several minutes and looks exactly like a hang
• --run= without a wandb account errors; WANDB_RUN=dummy + --run=dummy
• train bpb 0.721011 vs val 0.720161 — validation lower than train, which is what a genuine single-epoch run looks like (131 of 6,542 shards, no token seen twice)

Happy to share the full training log or eval CSV if useful.

---

Disclosure per repo AI policy: this post was drafted with LLM assistance. The run, configuration, logs and numbers are my own.