nanocoder-math

A from-scratch decoder-only language model for code and mathematics, trained under single-device compute constraints.

This project studies a practical question: how much capability can a small language model acquire when the entire pipeline — tokenizer, data, architecture, optimizer, and evaluation — must run on one machine? The codebase implements a complete GPT-style training stack in plain PyTorch, develops it on an Apple M4 (MPS) laptop, and runs unmodified on CUDA GPUs and CPU through a single hardware-abstraction layer. Rather than scale tokens or parameters, it pursues the levers available to a compute-limited researcher: data quality, optimizer efficiency, and architectural choices, each isolated through controlled ablations.

The repository began as a faithful reimplementation of the encoder–decoder Transformer of Vaswani et al. (2017), preserved in legacy_translation/, and was rebuilt into the decoder-only system described here.

---

Summary of results

All numbers below are measured on a single Apple M4 (10-core, 24 GB unified memory, MPS backend, fp32). They are reproducible with the commands in Reproducibility.

• Optimizer efficiency. Replacing AdamW with Muon (orthogonalized updates for 2-D weight matrices) lowers validation perplexity from 645 → 539 at identical token budget on the tiny model — a ~16% reduction in loss-per-token at no extra compute.
• Data-centric capability. Training the same model on a procedurally generated, formally verified arithmetic corpus yields 31.2% exact-match on held-out problems overall and 83.2% on single-digit arithmetic after only a few minutes of training, against ~0% for GSM8K/HumanEval at this scale.
• Hardware portability. A single precision policy selects bf16 + torch.compile
  • TF32 on modern CUDA, fp32 on Apple Silicon, and fp16+GradScaler on older CUDA, chosen at runtime from the detected architecture.
• Reinforcement learning from verifiable rewards. GRPO with an exact arithmetic verifier (no learned reward model) raises held-out exact-match from 53.8% → 57.4% overall and 93.0% → 98.4% on single-digit problems.
• Negative result, reported. An MLX reimplementation was benchmarked before any port and measured at 1.05× PyTorch-MPS throughput, below the 1.5× threshold set in advance; the port was therefore not undertaken.

Contributions

1. A compact, dependency-light decoder-only GPT with architecture choices exposed as ablation switches (positional encoding, normalization placement, feed-forward activation, QK-normalization, logit soft-capping, weight tying, initialization).
2. A streaming, budget-capped data pipeline producing uint16 memory-mapped token shards, with domain-mixing as a train-time parameter.
3. An implementation of the Muon optimizer compatible with the Metal backend, and a controlled comparison against AdamW.
4. A procedurally generated, verification-gated mathematics/code corpus and a matched exact-match evaluation, demonstrating the data-quality lever predicted by TinyStories and the Phi reports.
5. A reproducible cross-backend training system (MPS/CUDA/CPU) with auto-selected precision and torch.compile.

---

Method

Architecture

A standard pre-norm decoder-only Transformer (gpt/model.py). Defaults follow the GPT-2/nanoGPT line with rotary positional embeddings; each design axis is a configurable switch so it can be ablated independently.

Component	Default	Alternatives (ablation)
Positional encoding	RoPE	sinusoidal absolute
Normalization	pre-norm LayerNorm	post-norm
Feed-forward	SwiGLU	GELU MLP, squared-ReLU
Attention	scaled_dot_product_attention, causal	+ QK-RMSNorm
Output	tied embeddings, optional logit soft-cap	untied
Init	GPT-2 scaled residual	zero-init residual projections

Configuration	d_model	layers	heads	context	vocab	parameters (non-embedding)
tiny	384	6	6	1024	16k	16.8M (10.6M)
small	512	8	8	1024	16k	33.9M (25.7M)

Tokenizer

A byte-level BPE tokenizer (16k merges) trained on the target corpus (gpt/tokenizer.py). A small vocabulary is a deliberate choice: at this parameter scale a 50k-entry embedding table would dominate the budget, and a corpus-trained vocabulary compresses source code and LaTeX more effectively than a general-purpose tokenizer.

Data

Pretraining data is streamed from public, license-clean sources and capped by a token budget, avoiding multi-terabyte downloads (gpt/data/). Tokenized documents are stored as uint16 memory-mapped shards and read as random windows, giving O(1) access without loading the corpus into memory. Domain mixing (code vs. mathematics) is a sampling weight applied at batch construction, so mixture experiments require no re-tokenization.

Stage	Code	Mathematics
Pretraining	codeparrot/codeparrot-clean	open-web-math/open-web-math
Instruction tuning	CodeAlpaca, Python-instructions	GSM8K, MetaMathQA
Evaluation	HumanEval	GSM8K, synthetic arithmetic

A second, procedurally generated corpus (gpt/data/synth.py) supplies arithmetic, multi-step expressions, linear equations, GCD/LCM, fraction reduction, and templated Python functions. Every example is verified before inclusion — arithmetic and algebra via SymPy, code via execution — yielding an effectively unlimited, noise-free training signal generated entirely on-device.

Optimization

AdamW with decoupled weight decay and a linear-warmup cosine schedule is the baseline. The Muon optimizer (gpt/optim/muon.py) is provided as an alternative: it orthogonalizes the momentum update of each 2-D weight matrix through Newton–Schulz iteration, while embeddings, the output head, normalization parameters, and biases remain on AdamW. The orthogonalization runs in fp32 for numerical stability on the Metal backend.

Hardware abstraction

gpt/device.py is the single point at which hardware decisions are made. The backend is detected (mps > cuda > cpu) and --precision auto resolves to the appropriate policy:

Detected backend	Precision	Loss scaling	torch.compile	TF32
CUDA, bf16-capable (Ampere–Blackwell)	bf16	—	yes	yes
CUDA, older (sm_70/75)	fp16	GradScaler	yes	yes
Apple Silicon (MPS)	fp32	—	no	—
CPU	fp32	—	no	—

Two empirical findings on MPS motivated these defaults: fp16 autocast is slower than fp32 on Metal (no tensor-core speedup, added overhead), and non-blocking host-to-device copies race and corrupt batches. Both are handled in the code.

---

Experiments

Throughput (Apple M4, fp32)

Configuration	Context	Throughput	Tokens / 8 h
tiny (17M)	1024	4,585 tok/s	~132M
tiny (17M)	512	5,267 tok/s	~152M
small (34M)	1024	1,650 tok/s	~47M
small (34M)	512	1,060 tok/s	~31M

The tiny configuration attains a compute-optimal token-to-parameter ratio within a single overnight run; small is throughput-bound on this device and is intended for CUDA execution, where throughput is roughly 20–50× higher.

Optimizer ablation: Muon vs. AdamW

Identical model (tiny), data, seed, and token budget; only the optimizer differs.

Step	AdamW (val ppl)	Muon (val ppl)
50	1298	1066
100	796	680
150	645	539

Muon is uniformly ahead and ends ~16% lower in perplexity at equal tokens.

Data-centric result: verified synthetic mathematics

tiny, Muon + QK-norm, ~500 steps (minutes on the M4). Validation perplexity falls to 2.86 on the narrow corpus. Held-out exact-match accuracy:

Operand digits	Exact-match
1	83.2%
2	16.8%
3	5.6%
overall	31.2%

This is the central observation of the project: redirecting effort from scale to data quality converts a model that scores ~0% on standard benchmarks into one with a measurable, improving capability, consistent with the TinyStories and Phi findings.

Reinforcement learning from verifiable rewards (RLVR)

GRPO on the arithmetic task (gpt/rlvr.py): for each problem the policy samples a group of answers, each scored by the exact verifier (1 if the final number is correct), with group-relative advantages and a KL penalty to the pretrained reference. No learned reward model is required because the reward is exact.

Held-out exact-match	Before RLVR	After RLVR
1-digit	93.0%	98.4%
2-digit	14.4%	16.4%
overall	53.8%	57.4%

A modest but consistent gain from a few minutes of RL on the M4, demonstrating the verifiable-reward loop end-to-end at small scale.

Framework comparison: MLX vs. PyTorch-MPS

A FLOP-matched GPT was implemented in Apple's MLX and benchmarked against the PyTorch model on the same device (bench/mlx_gpt_bench.py):

Framework	Throughput (tiny, ctx 512)
MLX	9,363 tok/s
PyTorch-MPS	8,894 tok/s

The 1.05× ratio falls below the 1.5× threshold fixed in advance, so the port was not pursued. The benchmark is retained as a reproducible artifact of the decision.

Ablation suite

scripts/run_ablations.sh runs single-variable comparisons on tiny and emits overlaid loss curves: data mixture (70/30, 50/50, 30/70), RoPE vs. sinusoidal, pre- vs. post-norm, SwiGLU vs. GELU vs. squared-ReLU, AdamW vs. Muon, and the modded-nanoGPT switches (QK-norm, logit soft-cap, zero-init).

---

Evaluation protocol

• Language modeling. Per-domain validation perplexity and bits-per-byte (gpt/evaluate/perplexity.py).
• Mathematics. Exact-match on held-out, seed-disjoint synthetic problems, bucketed by operand magnitude (gpt/evaluate/arithmetic.py); GSM8K accuracy via numeric-answer extraction (gpt/evaluate/gsm8k.py).
• Code. HumanEval pass@1 with sandboxed, time-limited execution of generated programs (gpt/evaluate/humaneval.py).

Standard code/math benchmarks (GSM8K, HumanEval) are reported for completeness but are expected to be near zero at this scale; the synthetic exact-match metric is the instrument sensitive enough to track progress here.

Reproducibility

pip install -r requirements.txt

# pretraining corpus (streamed, budget-capped) and tokenizer
CODE_TOKENS=100000000 MATH_TOKENS=100000000 ./scripts/prepare_data.sh

# optimizer / architecture ablations on the tiny model
ABL_STEPS=600 ./scripts/run_ablations.sh

# narrow-domain (verified synthetic) track + arithmetic evaluation
./scripts/prepare_synth.sh && ./scripts/train_synth.sh

# RLVR (GRPO) on the verified arithmetic task, with before/after accuracy
BASE=checkpoints/synth_base.pt ./scripts/rlvr.sh

# overnight pretraining (Apple Silicon)
caffeinate -i ./scripts/pretrain.sh

# instruction tuning + full evaluation
BASE=checkpoints/tiny_base.pt ./scripts/finetune_eval.sh

# framework benchmark (reproduces the MLX decision)
python bench/mlx_gpt_bench.py

Running on CUDA (no code changes)

The 50-series (Blackwell, sm_120) requires a CUDA 12.8+ build; otherwise the same commands apply and precision/compilation switch automatically.

pip install torch --index-url https://download.pytorch.org/whl/cu128
python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_capability())"
python -m gpt.train --preset small --batch-size 48 --grad-accum 4 --max-steps 8000

Limitations and scope

This is a study of method under a hard compute budget, not a competitive model. At 17–34M parameters and well under a billion tokens, absolute performance on open-ended code and word-problem benchmarks is low by design. The contributions are the controlled comparisons, the portable training system, and the demonstration that data quality is the dominant lever at this scale. Reported numbers come from short runs on one device and should be read as relative, not as leaderboard results.

References

• Vaswani et al. Attention Is All You Need. NeurIPS 2017.
• Radford et al. Language Models are Unsupervised Multitask Learners (GPT-2). 2019.
• Su et al. RoFormer: Rotary Position Embedding. 2021.
• Shazeer. GLU Variants Improve Transformer (SwiGLU). 2020.
• Eldan & Li. TinyStories: How Small Can Language Models Be and Still Speak Coherent English? 2023.
• Gunasekar et al. Textbooks Are All You Need (Phi). 2023.
• Jordan et al. Muon and the modded-nanoGPT speedrun. 2024–2025.
• Karpathy. nanoGPT.