Supervised Pre-training Deep Dive

all_shards = sorted(shard_dir.glob("shard_*.npz"))
n_val = max(1, int(len(all_shards) * 0.05))   # ~15 shards from ~300
train_shards = all_shards[:-n_val]            # 95% for training
val_shards   = all_shards[-n_val:]            # 5% for validation

Shards are reserved at corpus load time and never trained on. With ~300 shards × 10K positions, that's ~150K positions exclusively for measuring generalization.

Teacher games are generated by scripts/generate_chess_corpus.py. Per game:

1. Random opening (4 plies). Without this, Stockfish-vs-Stockfish produces ~50 unique opening lines repeated thousands of times — useless diversity. Random first 4 plies give us ~20⁴ ≈ 160K distinct opening sequences from chess.Board().
2. Stockfish plays both sides with time_per_move=0.05s until the game ends.
3. Record every position as (encoded_state, move_played_idx, mover_player). We don't know the outcome until the game ends, so this accumulates in memory.
4. At game end, compute outcome via python-chess's board.outcome(claim_draw=True). The claim_draw=True is important — it makes the 50-move rule and 3-fold repetition terminate the game. Without it, games would only stop on checkmate, stalemate, insufficient material, or the unconditional 75-move/150-halfmove rule.
5. Backfill z for each recorded position from the final outcome × mover_player. Winner's positions get +1, loser's -1, draws 0.

Why Stockfish at 0.05s, not full strength

At 0.05s/move, Stockfish's effective ELO drops to ~2400 from its full ~3500. Three reasons we prefer it:

• Throughput. 0.05s/move × ~80 plies/game × 50,000 games / 10 workers ≈ 5 hours wall-clock. Full strength (say 2s/move) would be ~40 hours.
• Diversity. Faster play means more varied moves — Stockfish explores more lines without time to converge on one "best" move. Better for student generalization.
• Sufficient. ELO 2400 is vastly above the SP3 kill criterion (1500). No reason to over-pay.

Draws are expected and OK

Stockfish-vs-Stockfish at fast time controls produces ~60–70% draws. Both sides play accurately enough that ambiguous middle-game positions coast into 50-move-rule endgames. Typical corpus mix: 17% wins, 66% draws, 17% losses. That 33% decisive rate gives ~1M positions with z = ±1 across the 3M total — plenty of signal for the value head. The policy head benefits from all 3M positions regardless of outcome.

The corpus is stored as a series of .npz shard files, 10,000 positions per shard, written incrementally by each worker as its in-memory buffer fills. Each shard holds three numpy arrays:

states       : int8   (10000, 20, 8, 8)  ← encoded board planes
move_indices : int32  (10000,)            ← flat action indices [0, 4672)
outcomes     : int8   (10000,)            ← z values in {-1, 0, +1}

Why int8 and not float32?

The state encoding is nearly all 0s and 1s (piece presence, castling rights, etc.). As float32, a single shard would be 10000 × 20 × 64 × 4 = 5.1 MB. As int8: 1.3 MB. Compressed via np.savez_compressed, it shrinks further to ~190 KB per shard. For 300 shards (full corpus): ~60 MB on disk, vs ~1.5 GB if we'd stored as float32. Runtime cost is one cast int8 → float32 per batch when loading — negligible compared to forward+backward time.

Why 10,000 positions per shard?

A trade-off. Smaller shards → more file-system overhead, more np.load calls, but finer-grained progress. Bigger shards → fewer disk ops but more memory per load. 10K strikes the balance: each shard is ~190 KB compressed (fits in OS page cache), each worker accumulates one shard's worth in ~170 games (~20 minutes), and the loader processes one shard at a time with ~13 MB peak in-memory after the int8→float32 expansion.

Why incremental writes?

If we accumulated all positions per worker in RAM and wrote one big file at the end, a worker crash would lose everything. Flushing every 10K positions caps loss at the current buffer (~20 min of work). Plus the user sees visible progress (shard files appearing) instead of a 5-hour silent void.

Streaming with iter_batches

The pre-training loop never loads the whole corpus into memory. iter_batches is a Python generator (uses yield, not return) that yields batches one at a time, reading shards on demand:

def iter_batches(
    shards: list[Path],
    batch_size: int,
    device: torch.device,
    shuffle: bool,
) -> Iterator[tuple[Tensor, Tensor, Tensor]]:
    for shard_path in shards:
        data = np.load(shard_path)                              # one shard at a time
        n = data["states"].shape[0]
        order = torch.randperm(n).numpy() if shuffle else np.arange(n)   # per-shard shuffle
        for start in range(0, n, batch_size):
            idx = order[start : start + batch_size]
            s = torch.from_numpy(data["states"][idx]).to(torch.float32).to(device)
            m = torch.from_numpy(data["move_indices"][idx]).to(torch.int64).to(device)
            z = torch.from_numpy(data["outcomes"][idx]).to(torch.float32).to(device)
            yield s, m, z

• Stream, don't load. Peak memory stays bounded by one shard's worth (~13 MB float32) regardless of corpus size.
• Per-shard shuffle, not global. A true global shuffle across 300 shards would require either loading them all (defeats streaming) or complex multi-pass disk reads. Per-shard shuffle randomizes order within each shard's 10K positions; the caller can also shuffle the shard list between epochs for additional diversity.
• dtype casts at load. We saved int8/int32/int8 to minimize storage. At load we cast to what PyTorch expects: float32 for state features, int64 for class indices (cross-entropy expects int64), float32 for value targets.

The corpus generator spawns 10 worker processes (not threads) via multiprocessing.get_context("spawn"). Each worker:

• Owns its own Stockfish subprocess (configured Threads=1)
• Owns its own RNG (seeded with base_seed + worker_id)
• Plays games independently — no IPC with other workers
• Writes its own shards directly to disk (independent filenames)

Threads (Python) are limited by the GIL

CPython has a Global Interpreter Lock that allows only one thread to execute Python bytecode at a time per process. Fine for I/O-bound tasks where the GIL is released during waits. But our worker does Python-side work between Stockfish calls — encoding states, indexing moves, tracking buffers — and that IS bytecode execution. 10 threads in one process would serialize this through the GIL. Processes have no GIL issue: each is a fresh Python interpreter with its own memory space.

The "embarrassingly parallel" pattern

Corpus generation has zero inter-worker communication. No shared state, no locks, no merging. Each worker plays games, writes shards, exits. The orchestrator just spawns and joins them. The N-fold speedup is nearly linear up to physical core count. On a 10-physical-core machine, 10 workers achieve close to 10× the throughput of 1 worker. Hyperthreaded virtual cores give diminishing returns because hyperthreads on the same physical core share execution units.

Why we still use threads in parallel_selfplay.py

The chess SP3 parallel_selfplay.py (for AlphaZero self-play) does use a thread for the NN-server. That's because the NN-server needs to share GPU memory with the model — moving the model between processes is expensive (full state-dict copy) and fragile (spawn-context model loading has known PyTorch issues with MPS especially). The server thread releases the GIL during net(x) (PyTorch tensors are C++ under the hood), so multiple worker processes can submit requests concurrently without thread contention. So:

• CPU-bound Python work + no shared GPU → processes (corpus gen)
• GPU work + shared model state → thread inside the main process (NN-server in parallel_selfplay)
• CPU-bound Python work + needs GPU access → worker processes that send requests to a thread that owns the GPU (full parallel_selfplay architecture)

After pre-training completes, pretrained.pt has the exact same shape as the iter_NNNN.pt files AlphaZero training produces:

{
    "iteration":     0,                      # AZ refinement starts counting at iter 1
    "config":        <dataclass-as-dict>,    # informational only
    "best_net":      <state_dict>,           # ← this is what AZ refinement loads
    "candidate_net": <state_dict>,           # same as best_net (copy)
    "optimizer":     <optimizer state>,      # AdamW state from end of pretraining
    "_pretrain_epoch": N,                    # informational only
}

Format compatibility is non-negotiable — it's what lets Trainer.load_from_checkpoint(pretrained.pt) work without modification. The Trainer doesn't know or care whether the weights came from pre-training or from a prior AZ run; it just reads ckpt["best_net"] and proceeds.

Why candidate_net = copy(best_net)

Trainer maintains two copies: best_net (used to generate self-play games) and candidate_net (the one being trained). At the start of each iteration, gradient updates happen on candidate_net; periodically its weights are copied to best_net. By initializing both to the pretrained weights, the first AZ refinement iteration plays self-play with a strong (pretrained) network and trains on top of it.

Why iteration=0

AZ refinement starts at iter 1. If we saved iteration=N for some N > 0, the refinement loop would think it had already done N iterations and start at N+1. We set 0 so refinement counts from the beginning.

The integration test that protects this contract

tests/test_supervised.py::test_pretrained_checkpoint_loads_in_trainer_via_resume literally runs pretrain → save → Trainer.load_from_checkpoint → verify weights match. If anyone changes the checkpoint format on either side and breaks compatibility, this test fails. Regression guard for the entire bridge between Phase 2 and Phase 3.

Phase 1 — Corpus generation

python3 scripts/generate_chess_corpus.py --num-games 50000 --num-workers 10 --output data/chess_corpus/

Spawns 10 workers, each running Stockfish-vs-Stockfish games independently. Each worker writes its own shard files as buffers fill. Wall-clock: ~5–9 hours depending on machine. Output: ~300 shards totaling ~3M positions in data/chess_corpus/.

Phase 2 — Supervised pre-training

python3 -m alphazero pretrain --config configs/chess-pretrain.toml

Loads shards, streams batches, trains a 15M-param network for up to 3 epochs with held-out validation and early stopping. Wall-clock: ~30 min on M4 Pro MPS, ~2–4 hours on CPU. Output: pretrained.pt (Trainer-compatible checkpoint).

Phase 3 — AlphaZero refinement

python3 -m alphazero train --game chess --config configs/chess-refine.toml --resume-from pretrained.pt

Runs the existing AZ loop (now in parallel_selfplay.py with worker processes + NN-server thread) starting from the pretrained weights. Reduced parameters: 10 iterations × 100 games × 50 MCTS sims (vs SP3's 80/100/200). Wall-clock: ~10 hours. Output: checkpoints_refine/iter_NNNN.pt for each iteration.

Evaluation

python3 -m alphazero eval --game chess --checkpoint checkpoints_refine/iter_0010.pt --num-games 100

Plays 100 games vs Stockfish at ELO=1500. Pass criterion: win_rate >= 0.50.