d18-9tpp-hf-sft

nanochat miniseries

This repository is part of a miniseries of small (~360M–480M parameter) decoder-only transformers trained on top of Andrej Karpathy's nanochat codebase. The series varies three axes: depth (model size), tokens-per-parameter (pretraining horizon), and RoPE removal schedule (fraction of the pretraining token budget spent with RoPE before it is dropped for the remainder, used to study positional encoding in long-context generalization). A subset of the SFT models is additionally fine-tuned on a long-context mixture (_long variants).

All models share the same tokenizer: a BPE tokenizer with vocab size 32,768 trained on ~2B characters of the pretraining corpus.

Training pipeline

Each model goes through the following stages:

1. Tokenizer training — 32,768-vocab BPE trained on ~2B characters of the pretraining dataset.
2. Pretraining (base) — Next-token prediction on NVIDIA's ClimbMix-400B corpus, hosted at karpathy/climbmix-400b-shuffle. Horizon is controlled by target_param_data_ratio (aka "tpp" in model names), i.e. tokens trained per model parameter. Sequence length 4096, batch size 1,048,576 tokens, AdamW + Muon optimizer.
3. Supervised fine-tuning (SFT) — Instruction tuning on a mixture of:
   • HuggingFaceTB/smol-smoltalk — 460K general conversations
   • Synthetic identity conversations (from karpathy-public S3) — 1K rows × 2 epochs
   • cais/mmlu auxiliary_train — 100K rows × 3 epochs (multiple choice)
   • openai/gsm8k main — 8K rows × 4 epochs (math + tool use)
   • SimpleSpelling — 200K synthetic spelling examples
   • SpellingBee — 80K synthetic letter-counting examples
4. Long-context SFT (_long variants only) — Same mixture plus 100K rows of allenai/tulu-v2-sft-long-mixture, with sequence length extended to 8,192.

RoPE removal (drope) experiment

Model names containing drope_XX follow the recipe from "Extending the Context of Pretrained LLMs by Dropping Their Positional Embeddings": the model is pretrained normally with RoPE for the first XX% of its token budget, RoPE is then removed, and the remaining (100 − XX)% of the pretraining budget is used to recalibrate the model without positional encodings. For example, drope_50 means 50% of the token budget was spent with RoPE and the remaining 50% was spent with RoPE removed. This is intended to preserve the optimization benefits of RoPE early in training while producing a NoPE-style model that generalizes better to long contexts at inference time. Models without drope in the name keep RoPE in every layer for the full pretraining budget (theta = 100,000).

Model sizes

All models use head_dim=128, vocab=32,768, RMSNorm (ε=1e-6), SwiGLU MLP, and final logit softcapping at 15.0.

Released checkpoints

RoPE schedule column: none means RoPE is kept on for the full pretraining budget. A percentage (e.g. 50%) means RoPE is kept on for the first portion of the token budget and then removed for the remaining (100 − XX)% of pretraining, per the drope recipe above.

tpp = tokens-per-parameter pretraining horizon. Total pretraining token budgets:

drope variants use the same total token budget as their non-drope counterpart; the budget is split between the RoPE-on and RoPE-removed phases as described above.

Checkpoint format: which repo should I download?

For each model tag we publish four Hugging Face repositories:

• The base_checkpoints and chatsft_checkpoints artifacts are the raw nanochat outputs. They include the optimizer state (optim_*_rank0.pt) and metadata (meta_*.json with training config, val BPB, step number, etc.), so you can resume training or evaluate with the nanochat scripts exactly as produced by scripts.base_train and scripts.chat_sft.

• The hf_base and hf_sft artifacts are conversions of those same weights into the Hugging Face transformers layout (architecture name NanoChatForCausalLM, model_type nanochat). Load them with:

  from transformers import AutoModelForCausalLM, AutoTokenizer
  model = AutoModelForCausalLM.from_pretrained("crellis/nanochat-d20-20tpp-hf-sft", trust_remote_code=True)
  tokenizer = AutoTokenizer.from_pretrained("crellis/nanochat-d20-20tpp-hf-sft", trust_remote_code=True)
  

  use_rope in config.json reflects the drope setting: true for models that kept RoPE for the entire pretraining budget, and false for drope variants (where RoPE was removed partway through pretraining and the model was recalibrated without it). In the drope case, rotary embeddings are not applied at inference time.

Pick -hf-base / -hf-sft for inference. Pick -base / -sft only if you plan to continue training inside the nanochat codebase.

Inference sketch (HF format, SFT)

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

repo = "crellis/nanochat-d20-20tpp-hf-sft"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()

messages = [{"role": "user", "content": "Why is the sky blue?"}]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").cuda()
out = model.generate(inputs, max_new_tokens=256)
print(tok.decode(out[0][inputs.shape[1]:], skip_special_tokens=True))

Base (pretrained-only) checkpoints are next-token predictors and do not understand the chat template; use -hf-base for completion-style prompting and -hf-sft for chat.

Training compute

All runs were trained on a single H100 GPU via Slurm. Pretraining wall-clock ranges from ~4 hours (d18 @ 9tpp) to ~15 hours (d20 @ 40tpp); SFT adds ~30–90 minutes depending on variant.

Citation / acknowledgements

• Codebase: karpathy/nanochat
• Pretraining data: NVIDIA ClimbMix (via karpathy/climbmix-400b-shuffle)
• SFT data: HuggingFaceTB SmolTalk, CAIS MMLU, OpenAI GSM8K, AI2 Tulu-v2 long-mixture
• RoPE-removal recipe: Extending the Context of Pretrained LLMs by Dropping Their Positional Embeddings (arXiv:2512.12167)

License

MIT (inherits from the nanochat repository).