PAWN v2.0.0: PAWN as Shared-Weight Supernetwork; framework swap (PyTorch → JAX/Equinox + Optax)

[thomas-schweich]

Do not merge yet — per the user's directive, this PR is opened for review only. v2.0.0 is feature-complete on the integration branch.

What v2.0.0 is

PAWN 2.0.0 is a structural redesign on top of a complete framework swap:

1. Shared-weight supernetwork — one PAWNModel hosts three nested variants (small, base, large) that share parameters. Training the supernet trains all three variants jointly via a single multi-variant loss; the smaller variants are sliced views of the supernet's weights at eval/adapter time. validate_nested pins the nesting invariant at config-construction time.

2. All-JAX training + eval stack — Equinox modules + Optax optimizers + numpy + safetensors. No PyTorch in the training or eval paths. The two remaining torch touchpoints are pawn.torch_loader (thin loader for non-JAX consumers) and pawn._torch_legacy_fixture (frozen reference architecture used only by the legacy-converter parity tests). The CPU-jaxlib base install + the rocm/cu128/torch-loader extras keep the production surface torch-free unless you opt in.

3. Eight adapter strategies, one CLI — lora, film, unfreeze, bottleneck, hybrid, sparse, rosa (with three-phase schedule), specialized_clm (from-scratch). All dispatched via scripts/train_jax_adapter.py --strategy <name>.

4. Atomic checkpoint format with SHA-256 .complete sentinel — .tmp → rename → .bak sequence; interrupted overwrites always leave a recoverable checkpoint. Hashes are verified on every load.

5. Legacy-checkpoint converter — pawn.legacy.convert_legacy_checkpoint loads the three published PyTorch checkpoints (pawn-small, pawn-base, pawn-large) and produces JAX-format outputs at the same architecture (fp32 max-Δlogit ≈ 5×10⁻⁶ on the published-checkpoint test).

30 commits ahead of main. +22,500 / −50,300 across 181 files — the deletions are the entire pre-v2 PyTorch surface (lab, eval_suite, sweep, wandb_utils, cotrain, data, the old test tree, etc.).

Architecture map

	Pre-v2 (main)	v2.0.0 (jax_migration)
Framework	PyTorch + custom training loop	JAX/Equinox + Optax
Model	PAWNCLM per size	One supernet hosting nested small/base/large
Trainer	pawn.cotrain + scripts/train.py	pawn.trainer (K-step lax.scan) + scripts/train_jax.py
Adapter training	pawn.adapter_training	pawn.adapter_trainer (two-tier eqx.partition) + scripts/train_jax_adapter.py
Adapters	7 (no specialized_clm dispatch)	8, all dispatched via --strategy
Eval	pawn.eval_suite.* (PyTorch + polars + DataLoader)	pawn.{eval,probes,generation,lichess_eval} (numpy + Rust engine)
Generation	Recompute decoder	KV-cache decoder + variable-prefix-length grouping
Checkpoint	bare safetensors	safetensors + config.json + .complete sentinel (SHA-256)
Dashboard	Solara, default-on, hard dep	Solara, default-on but optional dashboard extra — base install ships without the visualisation stack
Sweep harness	Optuna + pawn.lab	gone (out of scope; expected to land in 2.x)
Pyproject deps	seaborn / matplotlib / polars / zstandard / solara / plotly / pandas / anywidget / optuna / fastmcp / starlette / wandb / ipykernel + torch	jax + equinox + optax + numpy + safetensors. torch and the dashboard stack are optional extras. uv.lock is substantially smaller.

Supernet + variants

Constant	d_model	n_layers	n_heads	d_ff
SUPERNET (= VARIANTS["large"])	640	10	10	2560
VARIANTS["base"]	512	8	8	2048
VARIANTS["small"]	256	8	4	1024
TINY_SUPERNET (verification scale)	192	4	3	768
TINY_VARIANTS["base"]	128	3	2	512
TINY_VARIANTS["small"]	64	2	1	256

All production variants share head_dim = 64. Legacy published pawn-large predates the supernet (head_dim = 80, n_heads = 8); the converter preserves its exact hyperparameters as a standalone ModelConfig that's not a nested slice.

Training stack

• K-step lax.scan in both pretraining and adapter trainers. K * B games per scan call; size K so per-step throughput is host-bound rather than launch-bound. State step is a jax.Array scalar (not Python int — the latter retriggers JIT recompile every step; pinned by tests).
• Multi-variant joint loss. Every pretraining step computes loss on the supernet plus each nested slice and sums them — the §5.3 supernet signal.
• Two-tier eqx.partition for adapter training. eqx.partition(model, adapter_filter(model)) produces a trainable subtree (adapters) and a frozen subtree (backbone); XLA dead-code-eliminates backbone gradients (~33% FLOP cut on the backward pass). Structural invariant — every array field of state.trainable.backbone is None after partitioning — is pinned across all 6 backbone-wrapping strategies by test_backbone_weights_are_frozen_all_strategies.
• Gradient clipping + warmup-cosine LR. optax.chain(clip_by_global_norm(1.0), adamw). decay_steps = total_steps (not total_steps - warmup — that double-subtracts; pinned).
• Padded-batch guards. A loss_mask-all-False chunk bypasses optimizer.update so AdamW's decoupled weight decay can't drift weights on zero-gradient padded steps.
• Factored embeddings. Each move token decomposes into src_embed[s] + dst_embed[d] + promo_embed[p] — ~14.8× fewer params on the input embedding table than a flat 1968 × d_model lookup.

Adapter strategies

Strategy	Module	Notes
lora	pawn.adapters.lora	LoRA + two-tier partition
film	pawn.adapters.film	per-layer γ⊙h+β; optional output FiLM (γ⊙logits+β)
unfreeze	pawn.adapters.unfreeze	per-layer gradient mask via optax.masked; the only strategy that uses gradient_mask
bottleneck	pawn.adapters.bottleneck	Houlsby MLP per layer
hybrid	pawn.adapters.hybrid	LoRA + FiLM composition
sparse	pawn.adapters.sparse	σ(s) / STE sparse mask
rosa	pawn.adapters.rosa	three-phase: LoRA warmup → gradient-magnitude mask gen → joint training. Phase gating swaps the adapter_filter, not via gradient_mask.
specialized_clm	pawn.adapters.specialized_clm	from-scratch standalone CLM (no backbone)

Eval suite

Surface	Module	Notes
Move-accuracy	pawn.eval	Per-phase breakdown. Argmax restricted to [0, NUM_ACTIONS); mask drops terminal move_{N-1} → PAD row (argmax can't reach PAD, otherwise every game biased accuracy by ~1/avg_game_length).
Linear probes (9)	pawn.probes	Streaming forward via lax.scan ys
Generation diagnostics (5)	pawn.generation	KV-cached + variable-prefix-length grouping. impossible_task_test / improbable_task_test require outcome_prefix_trained=True (otherwise return a {"_skipped": ...} sentinel — the conditioning signal only has interpretable meaning on prefix-trained models). zero_remaining_ply sizes the prefix to model.cfg.max_seq_len - 1 so the model genuinely has no room to generate.
Lichess Elo eval	pawn.lichess_eval	Atomic PGN cache (key includes max_games + _CACHE_VERSION); per-ply masking by side-to-move Elo so a 1500-vs-2500 game in the 1500 bucket scores only the 1500-player's plies.

Dashboard (pawn.dashboard)

The Solara metrics viewer is preserved as an optional extra: uv sync --extra dashboard installs solara / plotly / anywidget / pandas / starlette<1.0, and python -m pawn.dashboard --host 127.0.0.1 --port 8765 --log-dir logs reads each run's metrics.jsonl and renders charts. The Docker runtime / dev images bundle the extra by default; Caddy on port 8888 reverse-proxies the Solara server. Gated by PAWN_DASHBOARD (default on); the entrypoint checks python -c "import pawn.dashboard" before launching so a stripped-down install doesn't crash on startup.

Multi-perspective review — 8 rounds

After the feature work landed, the integration branch went through /subagent-review --pr 111 --auto --loop. Six lane agents (bug-detector / performance / type-correctness / test-risk / simplification / doc-accuracy) plus Codex ran in parallel each round; significant findings were fixed and the lanes re-ran until clean.

Round	Commit	Headline fix
1	7b9e09f	pawn.lichess_eval._cache_key omitted max_games; pawn.probes accidental perm_key = key overwrite; RoSA inactive-target 16 MB/forward zeros allocation; test_backbone_weights_are_frozen extended from LoRA-only to all 6 backbone-wrapping strategies; 5 CLAUDE.md doc errors.
2	79c9d59	lichess atomic-write tmp filenames could race-corrupt (fixed via pid+uuid); CPU-only skipif on KV-cache bit-equivalence tests; eval_jax upfront seq-len rejection test; lichess cache load-side guard test; 5 doc bugs.
3	bf97b88	pawn.eval move-accuracy was counting terminal move_N → PAD (Codex P2: biased every game by ~1/avg_game_length); try/finally cleanup on lichess tmp writes; Lichess-format PGN fixture so cache tests actually run; docs/ADAPTERS.md stale-legacy header.
4	76e6150	prepend_outcome=True loss_mask was dropping the outcome→m1 supervision step (Codex P2, re-raised); pinning test for the PAD-mask fix; pawn/eval.py docstring updated to reflect the new mask.
5	cf17ab5	deploy/entrypoint.sh still launched the deleted pawn.dashboard module (Codex P2 deployment break); upfront --val-every > 0 validation; _CACHE_VERSION bump prevents stale-cache silent reverts; prepend_outcome=True _materialise pinning test.
6	abc329e	filter_elo_slice was filtering whole games only — opponent's out-of-band moves contaminated reported per-Elo accuracy (Codex P2 third raise). Implemented per-ply masking by mover-Elo (depends on prepend_outcome parity). New pinning tests + val-every-zero test. Phase-4 cleanup: stale pyproject.toml dashboard exclude, deploy/vast.sh port-8888 forward.
7	3b1a1e3	Branch-coverage tests for the round-6 per-ply masking (side='white', side='black', prepend_outcome=True parity inversion).
post	7e44607 / 81d327b / ca1b49c	Gated impossible_task / improbable_task on outcome_prefix_trained; zero_remaining_ply now uses model.cfg.max_seq_len; fixed Rust engine parse_pgn_lichess initialising flat_tokens to 0i16 instead of PAD_TOKEN; H:MM:SS clock formatter in the lichess test fixture; restored the Solara dashboard as an optional dashboard extra (it wasn't actually PyTorch-coupled and the migration plan never argued for its removal — it just got swept up in the Phase 4 cleanup; the runtime image ships the extra by default and the entrypoint launches it under PAWN_DASHBOARD=1).

Verification

• uv run pyright pawn scripts tests — clean (modulo pre-existing scripts/vastai_score.py bs4 import; dashboard excluded per [tool.pyright].exclude).
• uv run --extra rocm pytest tests/ — 338 passed + 0 skipped (started at 317 + 1; +21 net tests from review rounds).
• Dashboard smoke: uv run --extra rocm --extra dashboard python -m pawn.dashboard --host 127.0.0.1 --port 18767 --log-dir /tmp — serves HTTP 200.
• End-to-end smoke verified:
  • Pretrain: 1K-step on TINY supernet, joint loss 22.41 → 17.46.
  • LoRA adapter: 500-step on TINY/base, val loss monotone; backbone bit-identical post-training.
  • Move-accuracy eval: converted pawn-small → 4.41% accuracy.
  • Linear probes on converted pawn-small: piece_type 0.61 → 0.83, side_to_move → 1.0 across layers.
  • Generation: KV-cache bitwise-equivalent to recompute (pinned on CPU).
  • All 8 adapter strategies dispatch end-to-end via the parameterised driver test.
  • RoSA three-phase: Phase 1 → 2 → 3 transition, step monotonicity, dense-mask Δ training.

Out of scope (deferred)

• Continuous LR across RoSA phases. Phase 3 replays the warmup ramp by design (Optax-internal step counter resets at the phase boundary; framework-level state.step preserved for metric monotonicity). Custom-schedule injection at the boundary is deferred.
• SIGTERM / graceful shutdown + HF-backed checkpoint push + W&B integration in scripts/train_jax{,_adapter}.py: not in this release.
• Sweep harness. pawn.lab + pawn.sweep + Optuna integration were removed. Expected to re-land in 2.x on top of the JAX trainer.
• pawn.eval_suite.bounds (theoretical accuracy ceiling) — not ported.
• Generation diagnostic zero_remaining_ply requires the corpus to be generated at corpus_max_ply >= model.cfg.max_seq_len - 1 to actually exercise the "no room to generate" semantic. The scenario silently skips on under-sized corpora rather than emitting misleading numbers; users running this on production-scale models should size --corpus-max-ply accordingly.
• Three engine PGN-buffer paths (parse_pgn_enriched, parse_pgn_lichess_filtered, uci_moves_to_tokens) share the same 0-init-then-PAD pattern that was fixed in parse_pgn_lichess, but none are called from the v2 Python surface. Fixed in this PR only on the path that's actually used.

Test plan

• Pyright clean across pawn / scripts / tests.
• Full pytest green (338 / 0 skipped) on CPU jaxlib.
• Engine builds (maturin develop --release) from a clean checkout.
• Dashboard installs + binds + serves under --extra dashboard.
• Smoke runs above all pass end-to-end.
• Reviewer: spot-check the three adapter strategies you care about most by running scripts/train_jax_adapter.py --strategy <name> --supernet tiny --variant base --total-steps 100 --k 50 --batch-size 2 --seq-len 16 --warmup-steps 5 --val-frac 0.1.
• Reviewer: convert one published checkpoint (uv run python scripts/convert_published_checkpoints.py) and run scripts/eval_jax.py against it to confirm move-accuracy lands in the ~4% band for pawn-small.

Per your directive, this PR remains unmerged pending your review.

[thomas-schweich]

Closing this PR — the branch is being abandoned in favour of a fresh from-scratch migration. See docs/jax_migration_plan.md on the new jax_migration branch for the replacement plan; this branch is preserved as ABANDONED-jax_migration on origin for reference.