chess-autocomplete-v1-checkpoints

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 repos (safetensors + ONNX), whose main loads through this project's load_model_repo. Each also carries an hf-transformers branch that AutoModelForCausalLM reads and a vllm branch for serving: chess-autocomplete-v1-91m, chess-autocomplete-v1-350m, chess-autocomplete-v1-700m-300k.

The models

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

350M

700M

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:

Compute slices (C = 6ND)

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

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

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.

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:

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.

The same surface fit on the pre-anneal points

Per-size fits against tokens

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

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.

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

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.