Skip to content

About

No description, website, or topics provided.

Resources

Stars

22 stars

Watchers

1 watching

Forks

Repository files navigation

How Model Growth, Recursion, and Boundary Operators Influence Scaling Exponents

The repository contains the code for the paper How Model Growth, Recursion, and Boundary Operators Influence Scaling Exponents by Zixi Chen, Akshay Vegesna, Samip Dahal, and Andrew Gordon Wilson.

Single-epoch and multi-epoch scaling comparisons

Figures · Training logs · Checkpoints

Setup

Use Python 3.10+ and H100 GPUs. The paper's standard ladders use one node with 8 GPUs; Untied-Grow d26 uses two nodes with 16 GPUs total.

uv sync --frozen
source .venv/bin/activate

Dependencies are pinned in uv.lock. FlashAttention-3 downloads its pinned kernel on first use. Training logs to W&B offline by default; no W&B account is needed.

Data

Prepare either corpus with the pinned dataset revision below. Each command writes train.pt, val.pt, and a provenance manifest, corpus.json, to the output directory. A 30B-token pool uses about 60 GB, plus space and host RAM for preparation.

# FineWeb
python prepare_corpus.py --corpus fineweb \
  --revision 9bb295ddab0e05d785b879661af7260fed5140fc \
  --out data/fineweb --train-tokens 30000000000

# FineWeb-Edu
python prepare_corpus.py --corpus fineweb-edu \
  --revision 87f09149ef4734204d70ed1d046ddc9ca3f2b8f9 \
  --out data/fineweb_edu --train-tokens 40000000000
Tokenization and validation splits

Both corpora use GPT-2 tokenization. FineWeb is the default corpus.

  • FineWeb: validation uses the first 10M tokens from sample-10BT. Training uses sample/100BT, excluding documents whose signature matches a validation document. The signature is a SHA-1 hash of up to 64 body tokens.
  • FineWeb-Edu: shards from sample/100BT are sorted by filename. The first shard supplies a 10.2M-token validation pool and is excluded from training; the remaining shards supply training tokens in order.

corpus.json records the resolved dataset revision, tokenizer and package versions, token and offset hashes, and shard provenance.

Reproduce the ladders

Run commands from the repository root. Each script in ladder_scripts/ trains a complete architecture ladder; bash -e stops on the first failed run.

# FineWeb (default dataset)
bash -e ladder_scripts/vanilla.sh

# FineWeb-Edu
export DATA_DIR=data/fineweb_edu
export RUNS_DIR=runs/fineweb_edu
export WANDB_GROUP=fineweb_edu
bash -e ladder_scripts/vanilla.sh
bash -e ladder_scripts/untied_grow.sh

DATA_DIR selects the folder containing train.pt and val.pt; RUNS_DIR selects where results are saved. Defaults are data/fineweb, runs, eight GPUs per node (NPROC_PER_NODE), and the W&B group scaling_ladders (WANDB_GROUP). Use a fresh RUNS_DIR when repeating an experiment.

Ladder scripts and recipes

Each standard ladder has 7 or 8 runs. TPP is the token-to-parameter ratio; d is the paper's depth coordinate.

Script Architecture Depth coordinates TPP LR exponent
vanilla.sh Vanilla 6–20, even 5 0.8
deep_vanilla.sh Deep Vanilla 6–18, even 6 0.7
operator_1.sh Operator-1 6–20, even 6 0.6
loop_2.sh Loop-2 6–18, even 6 0.6
untied_2.sh Untied-2 6–18, even 6 0.6
loop_grow.sh Loop-Grow 6–18, even 7 0.5
untied_grow.sh Untied-Grow 6–18, even 8 0.6
untied_grow_d26.sh Untied-Grow extrapolation 26 only 8 0.6

Width is 128d and the number of attention heads is d. The prelude/core/coda split divides d into thirds, with any remainder assigned to the core and then the coda. Deep Vanilla executes the same number of layers as Untied-2.

Learning rates follow LR(d) = 0.04 × (N(8)/N(d))^exponent, where N includes embedding and output-head parameters. Growth uses the initial K=2 parameter count. Fixed-depth runs use round(TPP × N / 524288) optimizer steps. Growth runs match the initial K=2 matrix-compute budget, switch from K=2 to K=4, and re-warm the learning rate for 40 steps. Final validation uses the final K.

All runs use bfloat16, a 2,048-token context, and 524,288 tokens per optimizer step. Exact settings and step counts are written in each script.

Untied-Grow d26 on FineWeb-Edu: two-node launch

The d26 run uses width 3,328, 26 heads, and an 8/9/9 prelude/core/coda split. At 77% of training it grows from K=2 to K=4, increasing active parameters from about 5.0B to 7.4B. All 7.4B parameters are allocated at initialization.

On both nodes, set the dataset and output paths:

export DATA_DIR=data/fineweb_edu
export RUNS_DIR=runs/fineweb_edu
export WANDB_GROUP=fineweb_edu

Replace node0 with node 0's reachable address. Use the same address and port on both nodes, and run the matching command on each:

# Node 0
NNODES=2 NODE_RANK=0 MASTER_ADDR=node0 MASTER_PORT=29500 NPROC_PER_NODE=8 \
  bash -e ladder_scripts/untied_grow_d26.sh

# Node 1
NNODES=2 NODE_RANK=1 MASTER_ADDR=node0 MASTER_PORT=29500 NPROC_PER_NODE=8 \
  bash -e ladder_scripts/untied_grow_d26.sh

Reproduce data-constrained scaling

The four scripts in data_constrained_scripts/ reproduce the scaling recipes in the third panel of Figure 5: Operator-1 and recurrence scaling, each with fixed or K=1-selected weight decay. They use a 100M-token FineWeb pool for 10 epochs.

Evaluate CORE accuracy and NLL

Evaluate a final checkpoint on one H100 GPU:

CUDA_VISIBLE_DEVICES=0 python eval.py \
  --checkpoint runs/vanilla_d08/checkpoints/final.pt

The default evaluation uses all 91,037 examples across 22 CORE tasks, with few-shot seeds 0, 1, 2. It downloads and verifies the benchmark bundle on first use; training data is not needed. Results are saved to the run's eval.json.

For a published checkpoint, download both final.pt and result.json, then pass their paths explicitly:

CUDA_VISIBLE_DEVICES=0 python eval.py \
  --checkpoint path/to/final.pt --result-json path/to/result.json \
  --out path/to/eval.json
Evaluation protocol and output metrics

Evaluation uses GPT-2 tokenization, the checkpoint's context length (2,048 in the paper), bfloat16, FlashAttention-3, and the final recurrence count (K=4 for growth runs). Seeds change the few-shot demonstrations; benchmark examples stay fixed. Zero-shot results are reused across seeds. Scores are averaged over seeds within each task, then equally across tasks.

JSON key Meaning
core_metric Full CORE accuracy: average of (accuracy − chance) / (1 − chance) over all 22 tasks; higher is better.
core_loss CORE answer NLL in nats/token: mean loss over the full correct-answer continuation, averaged over examples within each task, then equally over all 22 tasks; lower is better. Prompt tokens are excluded.
core_metric_filtered Mean centered accuracy on the paper's frozen 17-task subset. This excludes CommonsenseQA, BoolQ, and BIG-bench CS Algorithms, Language Identification, and Repeat Copy Logic.
core_metric_std Population standard deviation of full CORE accuracy across evaluation seeds.

Multiple-choice tasks select the candidate with the lowest mean token loss. Language-modeling accuracy requires every reference token to be predicted correctly. Answer NLL scores the full gold answer, including any prefix shared by the choices, and excludes prompt tokens.

Use --bundle-dir to select a benchmark directory (default eval_bundle/) and --out to select the results file. The evaluator checks the bundle against the paper's recorded content fingerprint.

Figures

Regenerate all 24 main-text and appendix figures from bundled data:

bash plotting_scripts/make_all.sh

See plotting instructions for dependencies, individual figures, and plotting your own runs. Outputs go to plotting_scripts/pdf/ and plotting_scripts/png/.

Public ladder runs

Click a final validation loss (lower is better) to open the run in loop-ladder-pub on W&B. Values are rounded to four decimals; rows use the depth coordinate d. A dash means no run at that depth.

FineWeb training logs — 58 runs across eight ladders
Depth Vanilla Deep Vanilla Deep Vanilla Grow Operator-1 Loop-2 Untied-2 Loop-Grow Untied-Grow
d6 3.5609 3.4500 3.4347 3.5075 3.4685 3.4286 3.4325 3.3551
d8 3.3296 3.2121 3.1963 3.2697 3.2323 3.1904 3.2002 3.1304
d10 3.1562 3.0394 3.0237 3.1016 3.0662 3.0190 3.0357 2.9625
d12 3.0292 2.9292 2.9140 2.9790 2.9473 2.9077 2.9184 2.8509
d14 2.9190 2.8194 2.8019 2.8663 2.8341 2.7940 2.8046 2.7416
d16 2.8259 2.7265 2.7077 2.7775 2.7427 2.7018 2.7154 2.6509
d18 2.7523 2.6625 2.6429 2.7044 2.6731 2.6378 2.6445 2.5870
d20 2.6834 — — 2.6373 — — — —

Deep Vanilla Grow is an additional paper baseline. The seven scripts in ladder_scripts/ cover the other columns.

FineWeb-Edu training logs — Vanilla and Untied-Grow, including d26
Depth Vanilla Untied-Grow
d6 3.3367 3.1348
d8 3.1048 2.9011
d10 2.9288 2.7345
d12 2.8035 2.6244
d14 2.6915 2.5110
d16 2.5989 2.4230
d18 2.5235 2.3592
d20 2.4529 —
d26 — 2.1372

Untied-Grow d26 is the 7.42B-parameter extrapolation run. See the two-node launch instructions to reproduce it.

Pretrained checkpoints

Click a full 22-task CORE score (core_metric, higher is better) to open its checkpoint on Hugging Face. Scores are the paper's mean chance-normalized accuracies, rounded to four decimals. A dash means no checkpoint is published.

FineWeb checkpoints — 51 models across seven ladders
Depth Vanilla Deep Vanilla Operator-1 Loop-2 Untied-2 Loop-Grow Untied-Grow
d6 0.0731 0.0683 0.0856 0.0761 0.0915 0.0938 0.1076
d8 0.0968 0.1016 0.1011 0.1198 0.1294 0.1284 0.1420
d10 0.1342 0.1454 0.1430 0.1482 0.1486 0.1565 0.1848
d12 0.1331 0.1855 0.1476 0.1612 0.1740 0.1855 0.1993
d14 0.1773 0.2127 0.1899 0.1981 0.2110 0.2074 0.2182
d16 0.2111 0.2348 0.2266 0.2132 0.2255 0.2388 0.2603
d18 0.2220 0.2469 0.2327 0.2425 0.2665 0.2568 0.2865
d20 0.2454 — 0.2478 — — — —
FineWeb-Edu checkpoints — Vanilla and Untied-Grow, including d26
Depth Vanilla Untied-Grow
d6 0.0799 0.1073
d8 0.0861 0.1523
d10 0.1344 0.1788
d12 0.1699 0.2103
d14 0.1753 0.2518
d16 0.2101 0.2871
d18 0.2378 0.3118
d20 0.2595 —
d26 — 0.3865

Acknowledgments and license

This project is released under the MIT License.

The optimizer builds on Nanochat and modded-nanogpt; the CORE evaluator is adapted from Nanochat. The license retains Andrej Karpathy's copyright notice for the Nanochat-derived portions. The prepared token-pool format was inspired by Slowrun.

About

No description, website, or topics provided.

Resources

Stars

22 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages