feat: framework swap to JAX/Equinox/Optax (v2 stack)

.gitignore

@@ -76,3 +76,6 @@ runs/
 
 # Codex review scratch file (ephemeral, sometimes left behind by `codex review`)
 .codex
+
+# Rust scratch binaries (e.g. `rustc -o rust_out`)
+/rust_out

DEFERRALS.md

@@ -0,0 +1,15 @@
+# Deferrals
+
+> Per `docs/jax_migration_plan.md` §9.4, a deferral is legitimate only
+> when actually doing the work is *nonsensical, impossible, or actively
+> detrimental*. Things on the §3 acceptance-criteria list cannot be
+> deferred under any reason.
+
+No deferrals.
+
+The framework swap implements the full §3 acceptance-criteria contract.
+Every v1 surface that existed on `main` has a v2 counterpart on this
+branch — including `scripts/benchmark.py`, which was rewritten against
+the JAX/Equinox/Optax stack (jit vs eager backbone steps, fresh-corpus
+vs pre-staged data-pipeline bench, multi-process JAX concurrency
+sweep, JAX adapter bench).

README.md

@@ -11,7 +11,7 @@ Feel free to use PAWN in your own experiments. PAWN is developed as a personal p
 
 ## Model Variants
 
-The model comes in three sizes, all trained from scratch on random chess games generated on-the-fly by a Rust-based chess backend. The v1.0.0 weights were trained together for 200K steps at batch size 256 on a single GPU — all three variants see the same random-game batches each step, with one forward/backward pass per variant in sequence (see [cotrain config](configs/cotrain_three_variants.json)). The numbers below come from the best 5K-cadence checkpoint by val loss (step 195,000 ≈ 49.9M sequences) for all three variants:
+The model comes in three sizes, all trained from scratch on random chess games generated on-the-fly by a Rust-based chess backend. The v1.0.0 weights were trained together for 200K steps at batch size 256 on a single GPU — all three variants see the same random-game batches each step, with one forward/backward pass per variant in sequence. **The v1 numbers below come from the best 5K-cadence checkpoint by val loss (step 195,000 ≈ 49.9M sequences) for all three variants under the PyTorch stack. The v2 stack (JAX/Equinox/Optax) replaces multi-variant cotraining with the supernet's joint loss; v2 weights publish to new HF repos (`pawn-{small,base,large}-v2`).**
 
 | Variant | d_model | Layers | Heads | Params | Top-1 | Legal rate | Game completion | Download |
 |---------|---------|--------|-------|--------|-------|------------|-----------------|----------|
@@ -47,28 +47,50 @@ uv sync --extra cu128   # NVIDIA GPU (or --extra rocm for AMD)
 Weights and data can be loaded directly from HuggingFace:
 
 ```bash
-uv run python scripts/train.py --run-type adapter --strategy bottleneck \
+uv run --extra rocm python scripts/train_jax_adapter.py \
+    --strategy bottleneck \
     --checkpoint thomas-schweich/pawn-base \
     --pgn thomas-schweich/pawn-lichess-full \
-    --bottleneck-dim 32 --lr 1e-4 --local-checkpoints
+    --bottleneck-dim 32 --lr 1e-4 --total-steps 200 --local-checkpoints
 ```
 
-### Pretrain from scratch
+The published `pawn-{small,base,large}` HF checkpoints are v1 PyTorch
+artifacts — they're loaded through `pawn.legacy.convert_legacy_checkpoint`
+which transposes weights into the v2 JAX layout. v2 republishes
+under new HF repos (`pawn-{small,base,large}-v2` or similar).
+
+### Pretrain the supernet
 
 Random games are generated on-the-fly; no dataset required:
 
 ```bash
-uv run python scripts/train.py --variant base --local-checkpoints
-
-# Or train all three variants simultaneously on shared data
-uv run python scripts/train.py --config configs/cotrain_three_variants.json
+uv run --extra rocm python scripts/train_jax.py \
+    --supernet base --total-steps 100000 --batch-size 256 \
+    --local-checkpoints
 ```
 
+The v1 cotrain path is GONE BY DESIGN — the supernet's joint loss
+(`sum` of per-variant cross-entropies on the same batch) replaces it.
+
 ### Run probes and diagnostics
 
 ```bash
-uv run python scripts/eval_probes.py --log-dir logs --device cuda
-uv run python -m pawn.dashboard --log-dir logs  # real-time monitoring
+# Move accuracy + per-phase
+uv run --extra rocm python scripts/eval_jax.py --checkpoint <converted-or-v2-dir>
+
+# Linear probes
+uv run --extra rocm python scripts/eval_probes_jax.py --checkpoint <converted>
+
+# 5 generation diagnostics (all gated on outcome_prefix_trained)
+uv run --extra rocm python scripts/eval_generation_jax.py \
+    --checkpoint <converted> --outcome-prefix-trained --edge-cases
+
+# Elo-stratified Lichess accuracy
+uv run --extra rocm python scripts/eval_vs_stockfish.py \
+    --checkpoint <converted> --pgn thomas-schweich/pawn-lichess-full
+
+# Real-time monitoring
+uv run --extra dashboard python -m pawn.dashboard --log-dir logs
 ```
 
 ## Architecture

deploy/pod.sh

@@ -326,8 +326,8 @@ cmd_launch() {
         echo "Usage: $0 launch <name> <command...>"
         echo ""
         echo "Examples:"
-        echo "  $0 launch exp1 scripts/train.py --variant base"
-        echo "  $0 launch exp1 scripts/train.py --run-type adapter --strategy bottleneck \\"
+        echo "  $0 launch exp1 scripts/train_jax.py --supernet base"
+        echo "  $0 launch exp1 scripts/train_jax_adapter.py --strategy bottleneck \\"
         echo "      --checkpoint thomas-schweich/pawn-base --pgn thomas-schweich/pawn-lichess-full \\"
         echo "      --elo-min 1800 --elo-max 1900 --bottleneck-dim 32"
         exit 1
@@ -389,7 +389,7 @@ case "${1:-}" in
         echo "  $0 create exp1 --gpu a5000"
         echo "  $0 create sweep1 --gpu a100-pcie --count 2 --community"
         echo "  $0 deploy exp1"
-        echo "  $0 launch exp1 scripts/train.py --variant base"
+        echo "  $0 launch exp1 scripts/train_jax.py --supernet base"
         echo "  $0 stop exp1"
         ;;
 esac

deploy/vast.sh

@@ -622,8 +622,8 @@ cmd_launch() {
         echo "Usage: $0 launch <name> <command...>"
         echo ""
         echo "Examples:"
-        echo "  $0 launch exp1 scripts/train.py --variant base"
-        echo "  $0 launch exp1 scripts/train.py --run-type adapter --strategy bottleneck \\"
+        echo "  $0 launch exp1 scripts/train_jax.py --supernet base"
+        echo "  $0 launch exp1 scripts/train_jax_adapter.py --strategy bottleneck \\"
         echo "      --checkpoint thomas-schweich/pawn-base --pgn thomas-schweich/pawn-lichess-full \\"
         echo "      --elo-min 1800 --elo-max 1900 --bottleneck-dim 32"
         exit 1
@@ -685,7 +685,7 @@ case "${1:-}" in
         echo "  $0 create exp1 --gpu 4090 --max-price 0.5"
         echo "  $0 create cheap1 --gpu 3090 --interruptible"
         echo "  $0 deploy exp1"
-        echo "  $0 launch exp1 scripts/train.py --variant base"
+        echo "  $0 launch exp1 scripts/train_jax.py --supernet base"
         echo "  $0 stop exp1"
         echo ""
         echo "Setup:"

docs/ACCURACY_CEILING.md

@@ -1,5 +1,12 @@
 # Theoretical Accuracy Ceiling
 
+> **v1 metrics disclaimer.** All metrics and benchmarks attached to
+> the published `pawn-{small, base, large}` HF repos are v1 PyTorch
+> numbers. v2 republishes to new HF repos
+> (`pawn-{small, base, large}-v2` or similar). The bridge that lets
+> v2 code load the v1 repos is `pawn.legacy.convert_legacy_checkpoint`.
+
+
 PAWN is trained on uniformly random chess games. At each position with
 N legal moves, the next move is drawn uniformly from those N moves, so
 the Bayes-optimal predictor (without outcome information) can do no

docs/ADAPTERS.md

@@ -1,5 +1,12 @@
 # Adapter Methods
 
+> **v1 metrics disclaimer.** All metrics and benchmarks attached to
+> the published `pawn-{small, base, large}` HF repos are v1 PyTorch
+> numbers. v2 republishes to new HF repos
+> (`pawn-{small, base, large}-v2` or similar). The bridge that lets
+> v2 code load the v1 repos is `pawn.legacy.convert_legacy_checkpoint`.
+
+
 PAWN is designed as a testbed for parameter-efficient fine-tuning. The frozen ~36M-parameter backbone provides learned chess representations from pretraining on random games; adapters specialize those representations for downstream tasks like predicting human moves at a given Elo level.
 
 All adapter implementations live in `pawn/adapters/`. Each wraps a frozen `PAWNCLM` backbone and exposes a uniform interface: `forward_hidden()`, `project_head()`, `forward()`, and `forward_generate()` (with KV-cache).

docs/ARCHITECTURE.md

@@ -1,5 +1,12 @@
 # PAWN Architecture
 
+> **v1 metrics disclaimer.** All metrics and benchmarks attached to
+> the published `pawn-{small, base, large}` HF repos are v1 PyTorch
+> numbers. v2 republishes to new HF repos
+> (`pawn-{small, base, large}-v2` or similar). The bridge that lets
+> v2 code load the v1 repos is `pawn.legacy.convert_legacy_checkpoint`.
+
+
 PAWN (Playstyle-Agnostic World-model Network for Chess) is a causal transformer trained on random chess games via next-token prediction. It learns chess rules, legal moves, and board state representations purely from move sequences, with no hand-crafted features or external game databases.
 
 This document describes the model architecture in detail.
@@ -82,7 +89,7 @@ x = x + FFN(RMSNorm(x))
 
 A final RMSNorm is applied after the last transformer block, before the output projection.
 
-**Attention.** Standard multi-head self-attention with no bias terms in any of the projection matrices (Q, K, V, output). Attention uses PyTorch's `scaled_dot_product_attention` with a causal mask combined with a padding mask. The padding mask ensures that PAD tokens are not attended to.
+**Attention.** Standard multi-head self-attention with no bias terms in any of the projection matrices (Q, K, V, output). The v2 stack (JAX/Equinox/Optax) materialises the attention `QK^T` scores plainly with a causal × padding mask combined via `jax.numpy.where`; the v1 PyTorch stack used `torch.nn.functional.scaled_dot_product_attention`. At seq 512 attention is ~12% of step FLOPs and plain attention sidesteps fused-kernel maturity under JAX-on-ROCm. The padding mask ensures that PAD tokens are not attended to.
 
 **Positional encoding.** [Rotary Position Embeddings (RoPE)](https://arxiv.org/abs/2104.09864) (Su et al., 2021) with base frequency 10000. RoPE is applied to the query and key vectors after projection, before the attention computation. Frequency tensors are precomputed for the full sequence length and stored as non-persistent buffers.
 
@@ -100,13 +107,23 @@ This uses three weight matrices per block instead of the standard two, with no b
 
 ## Model Variants
 
+**v1 (published HF checkpoints — PyTorch):**
+
 | Variant | d_model | Layers | Heads | Head dim | d_ff | Parameters |
 |---------|---------|--------|-------|----------|------|------------|
 | Small   | 256     | 8      | 4     | 64       | 1024 | 8.94M      |
 | Base    | 512     | 8      | 8     | 64       | 2048 | 34.65M     |
 | Large   | 640     | 10     | 8     | 80       | 2560 | 66.91M     |
 
-All variants use the same vocabulary, sequence length (512), and architectural choices. They differ only in width, depth, and head count. A `toy` variant (d=64, 2 layers, 4 heads) exists for testing.
+**v2 (JAX/Equinox/Optax — supernet + nested slices, all share depth + head_dim):**
+
+| Variant | d_model | Layers | Heads | Head dim | d_ff |
+|---------|---------|--------|-------|----------|------|
+| Small   | 256     | 10     | 4     | 64       | 1024 |
+| Base    | 512     | 10     | 8     | 64       | 2048 |
+| Large   | 640     | 10     | 10    | 64       | 2560 |
+
+The v2 stack pins `head_dim = 64` so width slices align to whole heads and RoPE is variant-invariant; all variants share the supernet's depth (10 layers) so the inner `[:d_V, :d_V]` of every weight matrix gives a valid sub-model. A `TINY_SUPERNET` (d=192, 4 layers, 3 heads) exists for verification runs that don't need production scale.
 
 The v1.0.0 parameter counts are slightly lower than the legacy `-legacy` repos with the same `d_model`/`n_layers`/`n_heads`, because the new 1,980-token vocabulary has roughly half the entries of the old 4,278-token vocab. The output projection (`lm_head`: `d_model → vocab_size`) is the only place vocab size enters the parameter count — factored input embeddings keep the input side compact regardless.
 

docs/TRAINING.md

@@ -1,5 +1,12 @@
 # Training Guide
 
+> **v1 metrics disclaimer.** All metrics and benchmarks attached to
+> the published `pawn-{small, base, large}` HF repos are v1 PyTorch
+> numbers. v2 republishes to new HF repos
+> (`pawn-{small, base, large}-v2` or similar). The bridge that lets
+> v2 code load the v1 repos is `pawn.legacy.convert_legacy_checkpoint`.
+
+
 ## Prerequisites
 
 - **Rust** (stable) -- required to build the chess engine native extension