NanoGPT Pro — Multi-Architecture NanoGPT Training & Evaluation Suite

Overview

NanoGPT Pro is a compiled Python package for training, running, and evaluating NanoGPT checkpoints across a broad zoo of attention and linear-attention architectures. It packages config-driven pretraining, dataset preparation scripts, and a standardized evaluation bridge so you can go from raw text to benchmarked language model on a single multi-GPU node.

The runtime ships as an audited, binary-only Linux CPython 3.14 wheel built on PyTorch 2.12. You do not need a source checkout for normal release use — install the wheel, point the bundled configs at your data, and train. Released alongside the wheel are tracked config/ training recipes, data/ preparation scripts for OpenWebText and FineWeb-EDU, and the NanoGPT Pro lm-evaluation-harness bridge.

Quick Start

Requires Linux x86_64, Python 3.14, and uv. For GPU training or evaluation, install a PyTorch build that matches your CUDA environment first. Download the nanogptpro wheel from the matching GitHub Release, then install it into a fresh environment:

curl -LO https://github.com/math-ai-org/nanogptpro/releases/download/v1.0.0/nanogptpro-1.0.0-cp314-cp314-linux_x86_64.whl

uv venv --python 3.14 .venv
source .venv/bin/activate
uv pip install ./nanogptpro-1.0.0-cp314-cp314-linux_x86_64.whl

python -c "import nanogptpro; print(nanogptpro.__version__)"
nanogptpro-generate --help

Train the reference gpt-mha-rope model on OpenWebText with a bundled config (downloads the dataset on first run):

export HF_HOME="${HF_HOME:-$PWD/.hf-cache}"

nanogptpro-train-openwebtext \
  config/train_gpt_mha_rope_small_adam_50BT_ctx1024_80g4.py \
  --data_path=data \
  --out_dir=out/gpt-mha-rope-openwebtext \
  --text_local_files_only=False \
  --wandb_log=False

Scale out across GPUs with torchrun, or swap in any other config/train_*.py recipe. Configs are plain Python — override any value from the command line with --name=value.

torchrun --standalone --nproc_per_node=4 \
  -m nanogptpro.train_adam_openwebtext \
  config/train_gpt_mha_rope_small_adam_50BT_ctx1024_80g4.py \
  --data_path=data --out_dir=out/gpt-mha-rope-openwebtext --wandb_log=False

Architecture Zoo

NanoGPT Pro is organized around Fast Weight Attention (FWA) — a unifying framework that casts recurrent state writes as online learning rules. DeltaNet, Mamba, Lightning Attention, and the project's own Falcon family are concrete implementations of this framework, alongside standard softmax-attention baselines. Every architecture is a drop-in config/train_*.py recipe sharing one training and evaluation pipeline.

Family	Implementations	Notes
Softmax Attention	MHA + RoPE, MHA + ALiBi, MHA + GRAPE, GQA, MQA, MLA, TPA	Causal transformer baselines spanning attention variants (GQA, MQA, MLA, TPA) and position encodings (RoPE, ALiBi, GRAPE).
FWA · DeltaNet	DeltaNet, Gated DeltaNet	Delta-rule fast-weight memory; chunked Triton and FLA kernels with context-dependent `beta`/`lambda`/`eta` gating.
FWA · Linear / SSM	Mamba 2, Mamba 3, Lightning Attention	Weight-decayed linear attention / state-space models with chunked Triton kernels and sliding states.
FWA · Falcon	Falcon-1/2/3, Falcon-1A/2A/3A	NLMS-normalized fast-weight updates: 1/2/3 are scalar, per-column, and sliding-window regression; the “A” suffix denotes the inner-product objective.

Features

Triton Kernels

Hand-written forward and backward Triton kernels for chunked linear attention and Fast Weight Attention, with fp32/bf16-safe numerics — no fp64 anywhere.

Architecture Zoo

One pipeline: softmax attention (MHA/RoPE, ALiBi, GRAPE, GQA, MQA, MLA, TPA) plus the Fast Weight Attention framework — DeltaNet, Gated DeltaNet, Mamba 2/3, Lightning Attention, and Falcon-1/2/3 & 1A/2A/3A.

Multi-GPU Training

Data-parallel pretraining via torchrun, tuned for 8×80G A100/H100 nodes with gradient accumulation and large global batches.

Large-Scale Data

Streaming and pre-tokenized pipelines for OpenWebText and FineWeb-EDU, plus arithmetic toy tasks for fast iteration.

Standardized Evaluation

A bundled LM Evaluation Harness bridge resolves NanoGPT Pro checkpoints automatically and benchmarks them on tasks like ARC and HellaSwag.

Compiled Release

An audited, binary-only Linux CPython 3.14 wheel with reproducible, config-driven runs and committed training recipes.

Citation

If you use NanoGPT Pro in research, please cite:

@misc{nanogptpro2026,
  title   = {NanoGPT Pro: A Multi-Architecture NanoGPT Training and Evaluation Suite},
  author  = {Team Math-AI},
  journal = {math-ai-org.github.io},
  year    = {2026},
  url     = {https://github.com/math-ai-org/nanogptpro}
}

NanoGPT Pro builds on nanoGPT and the EleutherAI LM Evaluation Harness.