This repository contains one chess-autocomplete model variant staged for inference.
Half the token budget of the other two releases. This is the annealed
300,000-step branch: it resumed the 700M trunk at step 240,000 and decayed the
learning rate to zero over the final 60,000 steps, so it saw 61.4B tokens.
The 91M and
350M releases each
ran the full 600,000 steps (122.9B tokens). The 700M trunk run stopped at 520,400
steps, so no 600k checkpoint of this size exists.
It is still the strongest of the three: 57.17% top-1 on the Allie 2022-blitz
test set, against 57.03% for the 350M and 56.07% for the 91M. Its own
token-scaling fit in the
checkpoints repo
puts it at ~57.35% had it seen the same 122.9B tokens.
Use TIME_CONTROL_MISSING_WORD and RATING_MISSING_WORD when metadata is not
available. The time-control token encodes one of four labels (bullet, blitz,
rapid, or classical), each with its own token. See
dataset.get_time_control_token.
The native PyTorch model returns logits over the output tokenizer vocabulary
(4135 ids). The ONNX artifacts wrap that model and return
bin_logits over raw 16-bit move words (65536 ids). These are different output
interfaces.
load_model_repo also takes a local path, so load_model_repo(".") works inside a
clone of this repo.
The PyTorch weights are stored in model.safetensors and loaded into the
transformers-native ChessTransformerForCausalLM, a standard causal-LM surface, so
loaded.model(input_ids=...).logits is all inference needs.
transformers and vLLM
main ships this project's own format, which transformers cannot load:
config.json uses the inference-config schema and model.safetensors carries the
training key names. Two branches repackage the same weights for the standard
loaders, with logits matching main to 0.0 absolute difference.
Its lm_head is zero-padded to vocab_size because vLLM rebuilds the head at
that width. Those padding classes sit at logit 0 and decode to no legal move,
so every request has to bias them out with SamplingParams(logit_bias={i: float("-inf") for i in range( config.num_move_classes, config.vocab_size)}). Each branch carries its own
README covering that contract and the input encoding, which is not the raw 16-bit
move words the ONNX artifacts take.
ONNX Runtime
python
1import numpy as np
2import onnxruntime as ort
34from chess_autocomplete import protocol
56session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])7bin_moves = np.asarray(8[[9 protocol.TIME_CONTROL_MISSING_WORD,10 protocol.RATING_MISSING_WORD,11 protocol.RATING_MISSING_WORD,12 protocol.GAME_START,13]],14 dtype=np.int32,15)16bin_logits = session.run(["bin_logits"],{"bin_moves": bin_moves})[0]
Four ONNX files are published. Every one takes a bin_moves input of shape
[batch, time]; they differ in how the weights are stored and in how much of the
output they return.
File
Weights on disk
Interface
bin_logits shape
model.onnx
FP32
bin_logits_v2
[batch, time, 65536]
model-bf16-fp32compute.onnx
BF16, cast to FP32 before every op
bin_logits_v2
[batch, time, 65536]
model-int8-blk128.onnx
8-bit block-wise MatMulNBits
bin_logits_v2
[batch, time, 65536]
model-int8-blk128-last.onnx
8-bit block-wise MatMulNBits
bin_logits_v1
[batch, 65536]
Read the interface from the ONNX metadata key chess_autocomplete.interface, or from
artifacts.*.interface in config.json. Do not infer it from the file name.
bin_logits_v2 returns logits for every position, so one forward pass scores a whole
game. bin_logits_v1 returns only the last position, which is what interactive play
wants: under v2 the per-move copy back from the device grows with the move number
(around 24 MB per move by ply 87), while v1 stays flat.
What the artifacts are for:
model.onnx is the FP32 compatibility artifact. Use it when you want FP32 weights
on disk.
model-bf16-fp32compute.onnx halves the weight size by storing floating weights as
BF16 and casting them back to FP32 before every op, so the graph computes and
outputs in FP32 and runs on any runtime with no BF16 operator support, including
onnxruntime-web (WebGPU/WASM).
model-int8-blk128.onnx is the smallest artifact. Linear weights are stored as
8-bit block-wise MatMulNBits (block size 128) and dequantized to FP32 at compute
time; the embedding and output head stay FP32 and activations are never quantized.
It is WebGPU-native (the MatMulNBits op runs on the onnxruntime-web WebGPU EP)
with no measurable strength loss (see Performance).
model-int8-blk128-last.onnx is the same quantized model under bin_logits_v1.
Its weights are verified byte-identical to model-int8-blk128.onnx at staging
time, so the INT8 performance row below covers both files. This is the one the
browser client downloads.
A model whose weights exceed protobuf's 2 GB message limit also ships an
external-data sidecar next to the graph (model.onnx.data beside model.onnx, and
likewise for the other artifacts). When a sidecar is present you must fetch both
files; ONNX Runtime resolves it by the relative path recorded in the graph.
The BF16 and INT8 artifacts are structurally checked before publishing and loaded
with ONNX Runtime CPU as a compatibility smoke test.
Performance
Held-out human-move-match on the ALLIE / Maia-3 Table 1 benchmark: top-1 move match
and legal-move NLL over the 2022-blitz test set, a clean training-excluded held-out
set. Each ONNX artifact is scored through the exact model it ships (the INT8 row is
the dequantized MatMulNBits weights, bit-faithful to the artifact), so these are the
numbers you get at inference. Δtop-1 is relative to the FP32 artifact.
Artifact
Precision
Size (MB)
Top-1 move match %
Δ Top-1 (pp)
NLL (legal)
Perplexity
model.onnx
fp32
2973
57.1686
0
1.29133
3.6376
model-bf16-fp32compute.onnx
bf16
1489
57.1656
-0.0029
1.29133
3.6376
model-int8-blk128.onnx
int8
807
57.1652
-0.0034
1.29136
3.6377
published references (matched)
MAIA-3-79M
57.1
MAIA-3-23M
56.6
MAIA-3-5M
55.4
ALLIE-ADAPTIVE-SEARCH
55.9
ALLIE-POLICY
55.7
MAIA-2
52.0
MAIA*
51.6
GPT-3.5
53.7
Third-party values are quoted from their source papers and carry about ±0.1 at 95%.
Their system names appear exactly as printed in Maia-3 Table 1, asterisk included.
The block-wise INT8 artifact is decision-equivalent to FP32 on this benchmark while
being the smallest download. Weight-only quantization keeps activations in FP32,
which avoids the accuracy collapse of dynamic (activation) INT8.
Converting Logits To Moves
The model predicts move tokens, not SAN strings. Do not take an unconstrained argmax
over the full vocabulary. Score the legal moves in the current board position and
choose from that legal set.
For PyTorch, logits are over the native output tokenizer vocabulary:
Call board.push(best_move) after selecting a move so the next prediction is decoded
against the updated legal move set.
Validation
Artifact
Validation
Status
Backend
Precision
Interface
Sample shape
model.safetensors
write
pass
safetensors.torch.save_file
model.safetensors
strict_load
pass
safetensors.torch.load_file
model.onnx
export
pass
torch.onnx
fp32
bin_logits_v2
[2, 2]
model.onnx
runtime
pass
onnxruntime.CPUExecutionProvider
fp32
bin_logits_v2
[2, 2]
model-bf16-fp32compute.onnx
export
pass
torch.onnx
bf16
bin_logits_v2
[2, 2]
model-bf16-fp32compute.onnx
onnx_checker_initializer_dtype_and_runtime
pass
onnx.checker+onnxruntime.CPUExecutionProvider
bf16
bin_logits_v2
[2, 2]
model-int8-blk128.onnx
quantize
pass
onnxruntime.MatMulNBitsQuantizer
int8
model-int8-blk128.onnx
onnx_checker_matmulnbits_and_runtime
pass
onnx.checker+onnxruntime.CPUExecutionProvider
int8
bin_logits_v2
[2, 2]
model-int8-blk128-last.onnx
quantize
pass
onnxruntime.MatMulNBitsQuantizer
int8
model-int8-blk128-last.onnx
onnx_checker_matmulnbits_and_runtime
pass
onnx.checker+onnxruntime.CPUExecutionProvider
int8
bin_logits_v1
[2, 2]
model-int8-blk128-last.onnx
weight_parity_with_model-int8-blk128.onnx
pass
onnx
int8
bin_logits_v1
Known Limitations
This model is trained for chess move autocomplete and is not a general chess engine.
It does not include Transformers AutoModel or trust_remote_code support.
Metadata-aware variants encode metadata as input tokens; no separate metadata tensor
path is supported. All four ONNX artifacts compute in FP32, and they differ only in
how weights are stored on disk and how much of the output they return.