Alfredvc/chess-autocomplete-v1-checkpoints: training checkpoints (with optimizer state)

34 checkpoints, 100 GB, across three model sizes and five token budgets. Each checkpoint is a torch.save payload carrying the full optimizer state (AdamW + Muon), the GradScaler and the dataset cursor, so you can resume training from it.

The models predict the move a human Lichess player of a given rating plays; Elo and time control are conditioning tokens. Weights-only inference repo (safetensors + ONNX, transformers/vLLM): chess-autocomplete-v1-91m, chess-autocomplete-v1-350m, chess-autocomplete-v1-700m-300k.

The models

repo dir	nickname	exact params	architecture	measured throughput
91m/	91M	91.4M	gpt2	421k tok/s (231 TFLOP/s)
350m/	350M	316.9M	gpt2_gpt2medium	136k tok/s (259 TFLOP/s)
700m/	700M	742.2M	gpt2_gpt2large	64k tok/s (283 TFLOP/s)

params is the exact count, embeddings included. The directory names are the GPT-2 nicknames we launched the runs under; every number below uses params.

Layout

Each size has one 600k-step trunk run plus five anneal branches. A branch for budget T resumes the trunk at 0.8·T and decays the LR to zero over the remaining 0.2·T steps; it differs from the trunk config only in max_iters and warmdown_iters. Each budget T therefore has two checkpoints:

• pre-anneal: the trunk checkpoint at 0.8·T, LR still hot. It lives under trunk/, and metrics.csv points at it there rather than duplicating it under anneal/.
• post-anneal: the branch's final checkpoint at T-1, LR decayed to zero.

91m/
  trunk/
    config.json
    training_config.yaml
    checkpoint-000015000.pt
    checkpoint-000030000.pt
    checkpoint-000060000.pt
    checkpoint-000120000.pt
    checkpoint-000240000.pt
    checkpoint-000480000.pt
    checkpoint-000599999.pt
  anneal/T018750/
    config.json
    training_config.yaml
    checkpoint-000018749.pt
  anneal/T037500/
    config.json
    training_config.yaml
    checkpoint-000037499.pt
  anneal/T075000/
    config.json
    training_config.yaml
    checkpoint-000074999.pt
  anneal/T150000/
    config.json
    training_config.yaml
    checkpoint-000149999.pt
  anneal/T300000/
    config.json
    training_config.yaml
    checkpoint-000299999.pt
350m/   (same layout)
700m/   (same layout; the trunk run stopped at 520,400 steps, so it has no
         480k or 600k checkpoint and no 600k budget row)

• config.json: the run's own config (model / train / hardware / tokenizers).
• training_config.yaml: the merged effective YAML, extends chain resolved, including data_config.shards. dataset_state.start_shard is a bare index into that list.
• metrics.csv / metrics.json: one row per checkpoint, with its evals, tokens, training_flops (6ND), gpu_hours, and its path in this repo.
• index.json: sha256 and byte size of every file.
• training-curves/: the full W&B history of every run here, one CSV each.
• figures/: the plots below, as PNG and SVG.

Evals

top1_move_match_pct on the Allie benchmark's 2022-blitz test set (884,049 positions, maia3_table1 surface). Δ acc = post-anneal − pre-anneal.

We hold the LR at max_lr for the first 80% of a budget, then anneal it to zero over the final 20%, as max_lr · (1 − √progress). The dashed line is the model before that anneal (the trunk at 0.8·T, LR still at max_lr); the solid line is the same run after it (the branch final at T, LR at zero).

91M

budget T (steps)	tokens	pre-anneal acc %	post-anneal acc %	Δ acc	post NLL	pre-anneal ckpt	post-anneal ckpt
18,750	3.8B	52.36	54.09	+1.73	1.4127	91m/trunk/checkpoint-000015000.pt	91m/anneal/T018750/checkpoint-000018749.pt
37,500	7.7B	53.01	54.72	+1.71	1.3873	91m/trunk/checkpoint-000030000.pt	91m/anneal/T037500/checkpoint-000037499.pt
75,000	15.4B	53.36	55.21	+1.85	1.3679	91m/trunk/checkpoint-000060000.pt	91m/anneal/T075000/checkpoint-000074999.pt
150,000	30.7B	53.73	55.56	+1.83	1.3531	91m/trunk/checkpoint-000120000.pt	91m/anneal/T150000/checkpoint-000149999.pt
300,000	61.4B	53.86	55.85	+1.98	1.3426	91m/trunk/checkpoint-000240000.pt	91m/anneal/T300000/checkpoint-000299999.pt
600,000 (trunk)	122.9B	53.96	56.05	+2.09	1.3347	91m/trunk/checkpoint-000480000.pt	91m/trunk/checkpoint-000599999.pt

350M

budget T (steps)	tokens	pre-anneal acc %	post-anneal acc %	Δ acc	post NLL	pre-anneal ckpt	post-anneal ckpt
18,750	3.8B	53.41	55.00	+1.59	1.3743	350m/trunk/checkpoint-000015000.pt	350m/anneal/T018750/checkpoint-000018749.pt
37,500	7.7B	54.01	55.64	+1.64	1.3499	350m/trunk/checkpoint-000030000.pt	350m/anneal/T037500/checkpoint-000037499.pt
75,000	15.4B	54.37	56.12	+1.75	1.3307	350m/trunk/checkpoint-000060000.pt	350m/anneal/T075000/checkpoint-000074999.pt
150,000	30.7B	54.65	56.49	+1.85	1.3171	350m/trunk/checkpoint-000120000.pt	350m/anneal/T150000/checkpoint-000149999.pt
300,000	61.4B	54.86	56.78	+1.91	1.3063	350m/trunk/checkpoint-000240000.pt	350m/anneal/T300000/checkpoint-000299999.pt
600,000 (trunk)	122.9B	55.00	57.01	+2.02	1.2985	350m/trunk/checkpoint-000480000.pt	350m/trunk/checkpoint-000599999.pt

700M

budget T (steps)	tokens	pre-anneal acc %	post-anneal acc %	Δ acc	post NLL	pre-anneal ckpt	post-anneal ckpt
18,750	3.8B	53.80	55.42	+1.62	1.3590	700m/trunk/checkpoint-000015000.pt	700m/anneal/T018750/checkpoint-000018749.pt
37,500	7.7B	54.43	56.06	+1.63	1.3348	700m/trunk/checkpoint-000030000.pt	700m/anneal/T037500/checkpoint-000037499.pt
75,000	15.4B	54.74	56.50	+1.77	1.3170	700m/trunk/checkpoint-000060000.pt	700m/anneal/T075000/checkpoint-000074999.pt
150,000	30.7B	55.03	56.88	+1.86	1.3031	700m/trunk/checkpoint-000120000.pt	700m/anneal/T150000/checkpoint-000149999.pt
300,000	61.4B	55.29	57.15	+1.86	1.2917	700m/trunk/checkpoint-000240000.pt	700m/anneal/T300000/checkpoint-000299999.pt

Architecture

All three share one decoder-only transformer: discrete move-token embedding, non-interleaved NeoX RoPE, SwiGLU MLP, LayerNorm, head dim 64, a 600-half-move context and a shared move+metadata vocabulary. They differ only in depth and width:

model	transformer blocks	model dim	attention heads
91M	12	768	12
350M	24	1024	16
700M	36	1280	20

Compute slices (C = 6ND)

Cooled (post-anneal) points grouped into bands whose C agrees within 20%.

compute C (6ND)	FLOP spread	sizes compared (accuracy %)	best
7.84e+18	1.15×	91M 55.21 · 350M 55.00	91M
1.61e+19	1.17×	91M 55.56 · 350M 55.64 · 700M 55.42	350M
3.23e+19	1.17×	91M 55.85 · 350M 56.12 · 700M 56.06	350M
6.46e+19	1.17×	91M 56.05 · 350M 56.49 · 700M 56.50	700M
1.26e+20	1.17×	350M 56.78 · 700M 56.88	700M
2.53e+20	1.17×	350M 57.01 · 700M 57.15	700M

The sweep grid is iso-token rather than iso-FLOP, so the bands span up to 1.17× in C (second column).

Fitted surface

L(N,D) = E + A·N^−α + B·D^−β, fit on the cooled points only, accuracy in error space (100 − acc):

accuracy = 58.91 − 16.7·(N/1e6)^−0.455 − 4.49·(D/1e9)^−0.385
R² = 0.99980

Minimised under 6ND = C: N* ∝ C^0.46, D* ∝ C^0.54. We fit β from five budgets per size and α from three sizes spanning 8×.

compute C	≈ H100-hours	optimal N*	optimal D*	D*/N*	predicted accuracy
3e+19	32	290M	17B	59	56.14%
1e+20	108	503M	33B	66	56.76%
3e+20	324	833M	60B	72	57.20%
1e+21	1,079	1447M	115B	80	57.58%
1e+22	10,793	4156M	401B	97	58.09%

The 91M trunk run sees 1,345 tokens per parameter; the fitted optimum at that compute is D*/N* ≈ 60.

Trained on the whole dataset

Alfredvc/chess-autocomplete-lichess holds 7.86B games = 528,651,645,230 move tokens (528.7B), which is 2.58M steps at 204,800 tokens per step. One pass, no repeats. We extrapolate each size below to that D twice, with two fits that extend along different axes:

• the L(N,D) surface, fit on all three sizes at once. It extends in N, so it gives a number for a size nobody trained (the frontier table above does exactly that).
• that size's own L(D) curve (E + B·(D/1e9)^−β, fit on its points alone). It extends in D, for that exact architecture, with no assumption tying it to the other sizes.

model	largest run	accuracy there	full dataset is	full dataset, L(N,D) surface	full dataset, that size's own L(D) curve	C = 6ND	GPU-hours
91M	122.9B tok	56.05%	4.3× that	56.36%	56.35%	2.9e+20	349
350M	122.9B tok	57.01%	4.3× that	57.29%	57.32%	1.0e+21	1,080
700M	61.4B tok	57.15%	8.6× that	57.68%	57.63%	2.4e+21	2,313

The largest run in this repo is 122.9B tokens, so you are reading both fits 4.3–8.6× beyond the tokens behind them. The surface's ceiling (N→∞, D→∞) is 58.91%.

The token count above covers the whole dataset. The training runs exclude the benchmark months (2019-12, 2023-12, all of 2022), which we do not subtract here.

The same compute, spent as the surface prefers:

a full-data epoch of	costs C	at that C, the surface's optimum is	its accuracy	vs the size itself on all the data
91M	2.9e+20	N* = 819M on D* = 59B tok	57.19%	56.36%
350M	1.0e+21	N* = 1451M on D* = 115B tok	57.58%	57.29%
700M	2.4e+21	N* = 2141M on D* = 183B tok	57.80%	57.68%

Wall-clock

Accuracy per GPU-hour, from each size's measured avg_tokens_per_second on the same GPU (1x NVIDIA H100 NVL). We interpolate accuracy along that size's own measured curve; a blank cell is a budget outside the range we ran.

GPU-hours (1× H100 NVL)	91M	350M	700M
5h	54.71% (8B tok)
10h	55.20% (15B tok)	55.22% (5B tok)
20h	55.55% (30B tok)	55.81% (10B tok)	55.58% (5B tok)
40h	55.84% (61B tok)	56.25% (20B tok)	56.17% (9B tok)
80h	56.05% (121B tok)	56.59% (39B tok)	56.60% (18B tok)
160h		56.86% (78B tok)	56.95% (37B tok)
240h		57.00% (118B tok)	57.11% (55B tok)

The same surface fit on the pre-anneal points

surface fit on	α (params)	β (tokens)	N* ∝ C^…	fitted accuracy ceiling	R²
post-anneal (cooled)	0.455	0.385	C^0.46	58.91%	0.99980
pre-anneal (hot LR)	0.520	0.545	C^0.51	56.37%	0.99831

Per-size fits against tokens

L(D) = E + B·(D/1e9)^(−β), one fit per size:

model	metric	fit (Chinchilla form)	R²
91M	final_val_loss	y = 1.33891 + 0.152229·(tokens/1e+09)^(-0.47)	0.99969
91M	mean_nll_legal	y = 1.31071 + 0.179692·(tokens/1e+09)^(-0.41974)	0.99995
91M	top1_move_match_pct	acc = 56.723 − 4.47137·(tokens/1e+09)^(-0.39468)	0.99997
350M	final_val_loss	y = 1.31226 + 0.143994·(tokens/1e+09)^(-0.46564)	0.99947
350M	mean_nll_legal	y = 1.27412 + 0.173892·(tokens/1e+09)^(-0.40902)	0.99995
350M	top1_move_match_pct	acc = 57.7379 − 4.56392·(tokens/1e+09)^(-0.3803)	0.99993
700M	final_val_loss	y = 1.29411 + 0.139766·(tokens/1e+09)^(-0.42448)	0.99945
700M	mean_nll_legal	y = 1.25521 + 0.171481·(tokens/1e+09)^(-0.3743)	0.99989
700M	top1_move_match_pct	acc = 57.9853 − 4.41413·(tokens/1e+09)^(-0.40372)	0.99971

Resuming training

import torch

ckpt = torch.load("91m/trunk/checkpoint-000240000.pt", map_location="cpu")
ckpt.keys()
# dict_keys(['model', 'optimizer', 'optimizer_muon', 'scaler',
#            'training_state', 'dataset_state'])

With the training code:

uv run python -m chess_autocomplete.pretrain \
    --config experiment-configs/final/pretrain_speedrun_weight_decay.yaml \
    --continue-from 91m/trunk --continue-from-iter 240000 --continue-id my-anneal \
    --out experiments/my-run

dataset_state is {start_shard, start_sample_idx}, an index into data_config.shards rather than a shard name, so it points at the same data only if the shard list is the one in the shipped training_config.yaml.

Metric definitions

• top1_move_match_pct: % of positions where the model's argmax move equals the human's.
• mean_nll_legal / perplexity_legal: NLL of the human move, renormalized over legal moves, on the Allie 2022-blitz test set.
• final_val_loss: next-token NLL on the held-out validation shard, from the training run's own final eval. A different distribution from mean_nll_legal. Pre-anneal rows have no value, since a mid-schedule checkpoint has no end-of-run eval block.
• training_flops = 6 · N · D, which leaves out the attention term (~3–4% at sequence length 200).
• All eval probability math runs in fp32, and we scored every point on one host with one code version.

Training data

Lichess standard rated games (lichess_db_standard_rated_*), most recent months first; data_config.shards in any training_config.yaml has the exact list and order. Sequence length 200 half-moves, batch size 1024, so tokens = iter × 204,800.

We benchmark these models against Maia-1/2/3 and Allie, so training excludes the months those test sets draw from: 2019-12, 2023-12, and all of 2022. The pipeline has no game-id holdout.

The full training corpus is at Alfredvc/chess-autocomplete-lichess; the evaluation sets (the Allie/Maia benchmark test data) are at Alfredvc/chess-autocomplete-eval-datasets.

Training curves

The tables above give endpoints. training-curves/ holds the runs behind them: every metric the training loop logged, for all 18 runs in this repo, as plain CSV. Nothing in the files is downsampled or smoothed, and the smoothing below happens only in the plots.

Each anneal branch resumes its trunk at 0.8·T and decays the LR to zero. Every endpoint in the next plot is a checkpoint you can download:

All three sizes run the same LR schedule, with Muon 5× hotter than AdamW. Below that, the gradient norms that schedule produces:

Throughput on one 1x NVIDIA H100 NVL. The wall-clock section above rests on these measurements. Bigger models use more of the GPU per step, so ranking the sizes by time and by FLOPs gives different answers:

What is in the CSVs

Two logging cadences share one table: eval metrics every eval_interval steps, the rest every log_interval steps (500 and 100 for these runs). A cell is empty where the run logged nothing for that metric at that step, so read every column against step rather than against a row index.

column	cadence	meaning
step	every logged step	optimizer step, absolute (a branch starts at 0.8·T)
tokens	derived	step × tokens_per_step, the x-axis of every card curve
wall_clock_s	every logged step	seconds since the run started
train_loss	eval	next-token NLL on a held-out slice of the training shards
val_loss	eval	next-token NLL on the held-out validation shard
valid_prob_mass	eval	probability mass the model puts on legal moves
game_end_accuracy	eval	accuracy on the game-terminating token
game_end_top1_accuracy	eval	top-1 accuracy on the game-terminating token
batch_train_loss	log	loss on the current microbatch, noisy by construction
lr	log	AdamW learning rate
muon_lr	log	Muon learning rate
unclipped_grad_norm	log	‖g‖ before grad_clip = 1.0 is applied
tokens_per_second	log	instantaneous, from the last iteration's duration
avg_tokens_per_second	log	averaged over the interval since the last log
samples_per_second	log	sequences/s, i.e. tokens_per_second / 200
t_flops	log	the run's own FLOP/s estimate (TiFLOP/s, 1024⁴ scaling)
shard_idx	log	index into data_config.shards, the dataset cursor
shard_sample_idx	log	position within that shard

step is absolute. An anneal branch's history starts at 0.8·T rather than at 0, because it continues its trunk.

run	steps logged	step range	tokens at the end	history
91M trunk	6,001	0 → 600,000	122.9B	training-curves/91m/trunk.csv
91M anneal T=18,750	38	15,100 → 18,750	3.8B	training-curves/91m/anneal-T018750.csv
91M anneal T=37,500	75	30,100 → 37,500	7.7B	training-curves/91m/anneal-T037500.csv
91M anneal T=75,000	150	60,100 → 75,000	15.4B	training-curves/91m/anneal-T075000.csv
91M anneal T=150,000	300	120,100 → 150,000	30.7B	training-curves/91m/anneal-T150000.csv
91M anneal T=300,000	600	240,100 → 300,000	61.4B	training-curves/91m/anneal-T300000.csv
350M trunk	6,001	0 → 600,000	122.9B	training-curves/350m/trunk.csv
350M anneal T=18,750	38	15,100 → 18,750	3.8B	training-curves/350m/anneal-T018750.csv
350M anneal T=37,500	75	30,100 → 37,500	7.7B	training-curves/350m/anneal-T037500.csv
350M anneal T=75,000	150	60,100 → 75,000	15.4B	training-curves/350m/anneal-T075000.csv
350M anneal T=150,000	300	120,100 → 150,000	30.7B	training-curves/350m/anneal-T150000.csv
350M anneal T=300,000	600	240,100 → 300,000	61.4B	training-curves/350m/anneal-T300000.csv
700M trunk	5,204	0 → 520,400	106.6B	training-curves/700m/trunk.csv
700M anneal T=18,750	38	15,100 → 18,750	3.8B	training-curves/700m/anneal-T018750.csv
700M anneal T=37,500	75	30,100 → 37,500	7.7B	training-curves/700m/anneal-T037500.csv
700M anneal T=75,000	150	60,100 → 75,000	15.4B	training-curves/700m/anneal-T075000.csv
700M anneal T=150,000	300	120,100 → 150,000	30.7B	training-curves/700m/anneal-T150000.csv
700M anneal T=300,000	600	240,100 → 300,000	61.4B	training-curves/700m/anneal-T300000.csv

Load them into your own W&B

training-curves/replay_to_wandb.py loads these CSVs into your own W&B project, so you can plot your run against ours on the same axes. It needs wandb and nothing else:

pip install wandb
python training-curves/replay_to_wandb.py --project my-project

# or just the runs you want ("91m" means every 91M run)
python training-curves/replay_to_wandb.py --project my-project --runs 91m/trunk 350m/anneal-T300000

# log locally now, upload whenever
WANDB_MODE=offline python training-curves/replay_to_wandb.py --project my-project
wandb sync --sync-all

The script keeps the original step numbers and loads each run's config.json as its W&B config, so the run-comparison filters work the way they do on a run you trained yourself. It logs tokens as a metric too; pick that as the x-axis to compare by data seen rather than by step.

License

apache-2.0. The Lichess data is CC0.