NanoGPT Pro — Multi-Architecture NanoGPT Training & Evaluation Suite

Overview

NanoGPT Pro is a public source and wheel release 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 public release source includes the full nanogptpro/ package, tracked config/ recipes, data/ preparation scripts for OpenWebText and FineWeb-EDU, and the NanoGPT Pro lm-evaluation-harness bridge. A compiled Linux CPython 3.14 wheel is also available for users who prefer a binary package.

Quick Start

Requires Python 3.14 and uv. For GPU training or evaluation, use a CUDA-capable environment with a matching PyTorch build. Clone the public release repository and install it from source:

git clone https://github.com/math-ai-org/nanogptpro.git
cd nanogptpro

uv venv --python 3.14 .venv
source .venv/bin/activate
uv pip install -e .

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

Train the reference gpt-mha-rope model on OpenWebText with a bundled config from the source checkout (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. The release wheel exposes the same command surface for users who prefer a compiled package.

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.

Source & Wheel Release

The public source checkout can run training and inference directly; an audited Linux CPython 3.14 wheel is available for compiled-runtime installs.

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.