Fully Nested Transformers

Code release for StairFormer, a fully nested Transformer architecture: a hierarchy of smaller Transformers nested inside a larger one, where every submodel's output is computed by a single forward pass of the full model.

The key ingredient is structure: block lower triangular weight matrices keep the residual stream of smaller models as a prefix of larger models', and PrefixRMSNorm keeps normalization from leaking information backward down the hierarchy. Together they make every layer filtration-preserving, so the top-left corner of the network is the smaller model.

• 📝 Blog post: Fully Nested Transformers
• 📄 Paper: OpenReview (ICML 2026 AdaptFM workshop)

This repo is a fork of karpathy/nanochat — the training harness, data pipeline, and evaluation are nanochat's, with the nested architecture and objectives added on top.

What's here

The nested-model additions, on top of stock nanochat:

File	What it does
nanochat/gpt_nested.py	The nested GPT: BlockTriangularLinear (only live weights stored), PrefixRMSNorm, per-block heads/attention, kfront submodel forward
nanochat/nested_objectives.py	Multi-budget training objectives — every submodel gets signal on every step
nanochat/prefix_ce_triton.py	Triton kernel for the prefix cross-entropy losses
nanochat/optim.py	Muon modified to orthogonalize block-triangular updates one row block at a time, so zero blocks stay zero
scripts/base_train_nested.py	Pretraining entry point for nested models
scripts/run_d24_sumflops_experiment.py	The blog's main experiment: dense baselines → summed-FLOP budget → nested run → CORE evals
tests/test_gpt_nested.py, tests/test_nested_objectives.py	Correctness tests, including exact submodel-equivalence checks

Setup

Dependencies are managed with uv:

uv sync --extra gpu    # CUDA (A100/H100/etc.)
source .venv/bin/activate

Then train the tokenizer and download data as in stock nanochat (see runs/speedrun.sh for the full pipeline).

Reproducing the main experiment

runs/stairformer_d24.sh runs the whole pipeline on an 8-GPU node: tokenizer + data, the dense d9 and d23 baselines at the default tokens-per-parameter ratio, then the d24 StairFormer — a depth-24 model (dim 1536, 24 heads) with two nested blocks, the inner model spanning the first 6 heads (dim 384) — trained to the sum of the two baselines' training FLOPs, followed by CORE evals of everything:

bash runs/stairformer_d24.sh

Or invoke the pieces directly. Train a nested model:

torchrun --standalone --nproc_per_node=8 -m scripts.base_train_nested -- \
  --depth=24 \
  --head-dim=64 \
  --nested-k=2 \
  --nested-block-head-counts=6-18 \
  --nested-loss-mode=big_ce_small_kl_pcgrad_hidden_manual_compiled \
  --kl-beta=0.15 --kl-tau=2.0 \
  --target-flops=3.68e19 \
  --device-batch-size=16

--nested-block-head-counts splits the attention heads into nested blocks (here 6 + 18 = 24); block boundaries coincide with head boundaries so attention stays filtration-preserving. See scripts/run_d24_sumflops_experiment.py for the full flag set used in the blog's run.

Evaluate any nested prefix with --nested-kfront (here the inner submodel), or omit it for the full model:

torchrun --standalone --nproc_per_node=8 -m scripts.base_eval -- --eval=core --nested-kfront=1

Results

One StairFormer training run (budget = sum of the two dense baselines' budgets) yields both models:

model	inference FLOPs/token	CORE ↑
dense small (d9)	1.52×10⁸	0.108
StairFormer small	1.86×10⁸	0.119
dense large (d23)	1.57×10⁹	0.262
StairFormer large	1.51×10⁹	0.246

Citation

@article{trost2026fullynested,
  title={Fully Nested Transformers},
  author={Trost, Avi},
  year={2026},
  month={July},
  url={https://avitrost.github.io/blog/fully-nested-transformers/}
}

Acknowledgements

Built on nanochat by Andrej Karpathy (MIT license). Trained on the ClimbMix dataset.