---
license: apache-2.0
library_name: chess-autocomplete
pipeline_tag: other
tags:
- chess
- pytorch
- training-checkpoints
- scaling-laws
---

# 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):
[`Alfredvc/chess-autocomplete-v1`](https://huggingface.co/Alfredvc/chess-autocomplete-v1).

![Accuracy against training tokens, pre-anneal and post-anneal](figures/anneal.png)

## 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; trunk still training, so no annealed 600k checkpoint yet)
```

* `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%.

![Accuracy against model size at near-fixed compute](figures/isoflop.png)

| 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-optimal frontier and the N*/D* split](figures/frontier.png)

| 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`](https://huggingface.co/datasets/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.

![Measured token curves extended to the full dataset](figures/projection.png)

| 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.

![Accuracy against wall-clock hours on one GPU](figures/isotime.png)

| 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

```python
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](https://github.com/Alfredvc/chess_autocomplete):

```bash
uv run python -m chess_autocomplete.pretrain_varlen \
    --config experiment-configs/final/pretrain_varlen_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`](https://huggingface.co/datasets/Alfredvc/chess-autocomplete-lichess); the evaluation
sets (the Allie/Maia benchmark test data) are at
[`Alfredvc/chess-autocomplete-eval-datasets`](https://huggingface.co/datasets/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.

![Validation loss over training](figures/training-loss.png)

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:

![Anneal branches leaving the trunk](figures/training-anneal.png)

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

![Learning-rate schedule](figures/training-lr.png)

![Gradient norm before clipping](figures/training-grad-norm.png)

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:

![Training throughput](figures/training-throughput.png)

### 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:

```bash
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.