cetusian/chess-sft-mix-200k

Chess SFT Mix — GM games + puzzles + Stockfish

A ~400,000-row supervised fine-tuning dataset assembled for competitive chess LLM training. Unified chat format, single system prompt, three complementary sources of signal.

This is the follow-up to `cetusian/chess-sft-lichess-2200`, which was pure GM behaviour cloning. This dataset adds tactical puzzles and engine-ground-truth labels on top of it, because BC alone caps at the teacher's (fallible) move quality.

Summary

All three components share a single chat format and a single system prompt, so the model learns one task: given a chess position, play the best move in SAN. The input shape tells the model whether the position came from a game (PGN) or a puzzle (FEN).

Format

Every row is an OpenAI-style messages list with exactly three messages:

{
  "messages": [
    {"role": "system",    "content": "You are a chess grandmaster. Given a chess position (as a PGN move sequence or a FEN), play the best move in standard algebraic notation (SAN)."},
    {"role": "user",      "content": "1. e4 e5 2. Nf3 Nc6 3. Bb5 a6 4. Ba4"},
    {"role": "assistant", "content": "Nf6"}
  ]
}

User message shapes:

• —PGN form (components A and C) — the move history as a standard PGN prefix. Example: "1. e4 e5 2. Nf3 Nc6 3. Bb5".
• —FEN form (component B) — "FEN: <fen>". Example: "FEN: r1bqk2r/pp1nbNp1/2p1p2p/...".

Assistant message: always a single move in SAN (e.g. "Nf6", "O-O", "e8=Q", "Qe6+", "Qxh7#").

This shape plugs directly into any modern SFT runner that accepts messages / conversation format — including TRL's SFTTrainer (with assistant_only_loss=True) and Surogate's type: conversation.

Component details

A. GM games (199,999 rows)

Imported verbatim from `cetusian/chess-sft-lichess-2200` — Lichess games with both players ≥ 2200 Elo, Termination == "Normal", ≥ 20 ply, FEN-deduplicated, opening-capped, colour-balanced. One random-ply position per game. The output is the move the human GM actually played.

Label quality (from the upstream card's Stockfish probe, depth 18, 2k positions): 41.1% top-1 agreement, 68.3% top-3, median 13cp loss, 7.5% blunder rate.

B. Lichess puzzles (150,000 rows)

Sampled from `Lichess/chess-puzzles` (5.8M puzzles, CC0) with uniform stratification across five rating buckets:

For each puzzle we:

1. 1.Start from the given FEN (position before the opponent's setup move).
2. 2.Apply the first UCI move of Moves — the opponent's move that sets up the tactic.
3. 3.Convert the second UCI move — the player's first correct answer — to SAN.
4. 4.Emit {user: "FEN: <position-after-setup>", assistant: "<SAN-answer>"}.

FENs are deduplicated across puzzles. Puzzles are where most of the dataset's # (mate) labels come from — GM games almost never record the final mating move.

C. Stockfish-enriched GM positions (50,000 rows)

Fifty thousand positions resampled from component A's PGN prefixes, analysed with Stockfish 16 at depth 15 on 64 parallel workers. The assistant output for these rows is Stockfish's preferred move, not the GM's. When SF agrees with the GM these rows reinforce; when it disagrees (≈59% of the time per the probe in A) the model sees a contrasting label for a familiar-looking position and learns the engine's preference.

Generation pipeline: data/prepare_mixed_sft.py (and companion data/stockfish_probe.py) in the source repository.

Statistics

Size & composition

Output shape

The tactical footprint is much heavier than in component A alone — 7× checks and ~11k mates (vs zero), courtesy of the puzzle mix.

User-message length (characters)

FEN inputs are short and fixed-ish (~70 chars); PGN prefixes are longer. The long tail still comes from deep endgame PGNs in the GM component.

Usage

Load

from datasets import load_dataset
ds = load_dataset("cetusian/chess-sft-mix-200k")
# DatasetDict({
#   train: Dataset(features=['messages'], num_rows=391999),
#   validation: Dataset(features=['messages'], num_rows=8000)
# })

Train (TRL SFTTrainer)

from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig

MODEL = "Qwen/Qwen3-1.7B"    # or Qwen3-8B, Qwen3.5, SmolLM2, etc.
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype="auto")

trainer = SFTTrainer(
    model=model, tokenizer=tok,
    train_dataset=ds["train"], eval_dataset=ds["validation"],
    args=SFTConfig(
        output_dir="runs/chess-sft-mix",
        max_seq_length=1024,            # covers >99% of rows
        per_device_train_batch_size=16,
        gradient_accumulation_steps=2,
        num_train_epochs=1,             # 392k rows is ample
        learning_rate=2e-5,
        assistant_only_loss=True,       # CRITICAL: supervise only the 2–8 move tokens
        packing=False,                  # keep position boundaries
        bf16=True, warmup_ratio=0.03,
    ),
)
trainer.train()

Train (Surogate)

datasets:
  - path: cetusian/chess-sft-mix-200k
    split: train
    type: conversation
    messages_field: messages

validation_datasets:
  - path: cetusian/chess-sft-mix-200k
    split: validation
    type: conversation
    messages_field: messages

Design notes

• —Why one system prompt, not four per task? The user-message shape (PGN vs FEN-tagged) already tells the model which mode it's in. Multiple prompts would just over-specialise per task and force extra mode-routing for no benefit.
• —Why 150k puzzles out of 5.8M available? Uniform 30k-per-bucket sampling gives the model tactics at every skill level without letting the huge beginner/intermediate bucket drown positional play signal.
• —Why only 50k Stockfish rows? Depth-15 engine labels are the most expensive component. 50k is the sweet spot where correctness signal is visible in loss curves without blowing past a few minutes of 64-core compute.
• —Why no tablebase? Syzygy 6-piece is ~150 GB on disk; marginal signal versus the Stockfish component on the same positions. Easy to add later.

Limitations

• —Mixed-source label contradictions. Some positions appear in both A (with GM label) and C (with SF label). When the two disagree the model sees both moves as correct — acceptable for a behaviour-cloning base but something to be aware of.
• —Puzzles drop opening context. Component B uses FEN: only, so the model does not see how the puzzle was reached. This is intentional — puzzles are about the immediate tactic — but means the FEN and PGN modes train somewhat separately.
• —Single-move-per-row puzzles. Only the player's first move in each puzzle is emitted. Subsequent forced moves are dropped to keep row count bounded and to avoid upweighting long puzzles.
• —Ceiling is still engine-level, not world-championship. Stockfish at depth 15 is very strong but not tablebase-perfect. The Stockfish component's labels inherit this cap.

Source & license

• —Lichess games: `Lichess/standard-chess-games`, CC0.
• —Lichess puzzles: `Lichess/chess-puzzles`, CC0.
• —Stockfish evaluations: Stockfish 16, GPLv3 (evaluations are not copyrightable; the engine is only used as a tool).

This derivative dataset is released under CC0 1.0 Universal.

Citation

@misc{lichess-database,
  title  = {Lichess Open Database},
  author = {{Lichess}},
  year   = {2024},
  url    = {https://database.lichess.org/},
  note   = {CC0 1.0 Universal}
}
@misc{stockfish,
  title  = {Stockfish: A strong open source chess engine},
  author = {The Stockfish developers},
  url    = {https://stockfishchess.org/},
  note   = {GPLv3}
}