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

[thomas-schweich]

Summary

Full framework swap from PyTorch to JAX/Equinox/Optax per the
from-scratch parity-bar plan in docs/jax_migration_plan.md. Every
v1 user-facing surface is preserved with a v2 counterpart; the v1
GONE-BY-DESIGN list in plan §6 is fully retired.

The migration is structured as 16 sections (S1–S16) — see commit log
for the section-by-section squash history.

§3 acceptance criteria

All 20 criteria pass on this branch. See final_smoke.md at the repo
root for the verification command + pasted output for each one.
Highlights of the live runs in that artifact:

• §3.4: scripts/convert_published_checkpoints.py runs through
  thomas-schweich/pawn-{small,base,large} and produces three v2 JAX
  checkpoints. The synthetic-v1 parity test asserts the converter is
  numerically within mean Δlogit ≤ 1e-3 / max ≤ 1e-4.
• §3.6: 1000-step train_jax.py --supernet tiny --total-steps 1000 --batch-size 16 --seq-len 64 --k 50 — loss decreases from 19.7 →
  17.5, no NaNs, .complete-sealed checkpoint at step 1000.
• §3.7: 200-step LoRA on v1-converted pawn-base — loss 3.42 → 3.30
  with the held-out validation split as val source.
• §3.9: eval_jax.py against converted pawn-base reports top-1
  accuracy 8.72%, within 0.15pp of v1's reported 8.57%.
• §3.14: 3-trial Optuna LoRA sweep, exits 0, best trial is
  {rank=4, targets='qkv', lr=0.0056} with val_loss=3.27.

Test suite

uv run --extra rocm --extra wandb --extra dashboard --extra lab pytest tests/ -q -m "not slow" → 619 passed, 0 failed.

uv run --extra rocm pyright → 0 errors across pawn/,
scripts/, tests/.

Deferrals

None. DEFERRALS.md at the repo root is empty — every §6 surface
ports forward. See the file for the explicit "No deferrals"
declaration.

What lands here

• pawn/model.py — Equinox PAWNModel supernet (lax.scan over
  n_layers, plain attention, RoPE, SwiGLU, factored embeddings).
• pawn/trainer.py — pretraining with K-step lax.scan inside one
  @eqx.filter_jit body; supernet joint loss = sum of per-variant
  cross-entropies on the same batch.
• pawn/adapter_trainer.py — two-tier eqx.partition frozen-
  backbone / trainable-adapter training; 10-key STRATEGIES
  dispatch table including the three RoSA modes.
• pawn/adapters/{lora,film,bottleneck,hybrid,sparse,rosa,unfreeze, specialized_clm}.py — full v1 adapter surface ported.
• pawn/eval.py + pawn/probes.py + pawn/generation.py +
  pawn/lichess_eval.py + pawn/eval_suite/diagnostics.py — eval
  surface incl. 5 outcome-gated generation diagnostics + edge-case
  Rust-engine integration.
• pawn/legacy.py + pawn/_torch_legacy_fixture.py — single bridge
  from v1 PyTorch HF checkpoints to v2 JAX safetensors, with a
  parity test against a synthetic v1 fixture.
• pawn/lifecycle.py — HFPushTracker (async ThreadPoolExecutor),
  SIGTERM handler, load_resume_state with step-splicing.
• pawn/run_config.py — pydantic BaseRunConfig /
  PretrainConfig / AdapterConfig / SpecializedCLMConfig /
  RunConfig discriminated union, all extra="forbid".
• pawn/sweep.py + pawn/lab/ + pawn/dashboard/ — Optuna driver,
  FastMCP lab server, Solara dashboard (all torch-free).
• pawn/_sentinel.py — stdlib-only .complete SHA-256 sentinel
  helpers used by every save/load path.
• pawn/corpus.py + pawn/lichess_data.py — Rust-engine random
  corpus + Lichess parquet → JAX Corpus, on-disk cache, multi-
  epoch tiling.
• scripts/train_jax.py + scripts/train_jax_adapter.py +
  scripts/eval_jax.py + scripts/eval_probes_jax.py +
  scripts/eval_generation_jax.py + scripts/eval_vs_stockfish.py
  • scripts/sweep.py + scripts/convert_published_checkpoints.py
  • scripts/run_evals_backbone.py — v2 entry points.
• scripts/benchmark.py — fully rewritten against JAX/Equinox/Optax
  (jit vs eager backbone + adapter steps, fresh-corpus-per-step vs
  pre-staged data-pipeline bench, multi-process JAX concurrency
  sweep).
• engine/src/lib.rs — PAD-init at four PGN-token-init sites so a
  0-initialised tail no longer aliases a legal move token.
• tests/test_jax_*.py — JAX-side test per module plus
  tests/test_public_api.py pinning the post-swap public surface.

What goes away (GONE BY DESIGN per §6)

• pawn/cotrain.py → replaced by pawn.trainer.supernet_joint_loss.
• pawn/gpu.py → JAX handles device config; PAWN_ALLOW_CPU=1
  preserved as the operator escape hatch in both v2 entry points.
• pawn/data.py + pawn/data_utils.py → pawn.corpus + pawn.lichess_data.
• pawn/adapter_training.py → pawn.adapter_trainer.
• pawn/specialized_clm.py → pawn.adapters.specialized_clm.
• pawn/lichess_cache.py → pawn.lichess_data's on-disk cache.
• pawn/eval_suite/{worker,generation,lichess,probes}.py →
  pawn.{generation,lichess_eval,probes}.
• pawn/lab/sweep.py → pawn.sweep.
• scripts/{train,eval_accuracy,eval_probes,export_hf_repo}.py →
  scripts/*_jax*.py.

Multi-lane review

The S16 multi-lane review pass surfaced one Critical pile in the
adapter trainer (saving the wrong tensor, the unfreeze no-op, the
unused FiLM beta, the v1-vs-v2 kaiming bound divergence, the
generate_corpus-per-step performance regression) plus assorted
type / doc / dead-code findings — every item is addressed in the
S16-review fix commit (6bf5dc7). The fifth-pass spec-alignment
gate returned PASS.

Test plan

• CI: uv sync --extra rocm exits 0 (criterion 1)
• CI: uv run --extra rocm pytest tests/ is green (criterion 3)
• CI: uv run --extra rocm pyright reports 0 errors
• Manual: rerun a full set of §3 verification commands on a
  ROCm pod / vast.ai instance; compare against final_smoke.md
• Manual: convert published v1 HF checkpoints (pawn-{small,base, large}) and verify the v2 forward path produces accuracy
  within ±0.5pp of v1's published numbers
• Manual: publish v2 supernet-derived checkpoints to
  thomas-schweich/pawn-{small,base,large}-v2

Plan: docs/jax_migration_plan.md

[claude]

Optimizer state is always cold-restarted on resume.

opt_state is reinitialised from scratch here even though save_model accepts (and stores) an optimizer_state sidecar. The code that actually calls save_model in scripts/train_jax.py never passes optimizer_state, so there is nothing to load — but the result is that every --resume silently resets Adam's first/second moment estimates to zero.

For short fine-tuning runs this is usually harmless. For long pretraining runs resumed mid-cosine (where momentum has converged) the cold restart can cause a loss spike and a multi-hundred-step recovery period.

To preserve the v1 behaviour, the trainer should:

1. Serialise the opt state via optax.tree_utils.tree_map_params + save_file and pass it to save_model(…, optimizer_state=…).
2. Load it here via st_load + reconstruct with eqx.apply_updates style.

If cold restart is intentional for simplicity, add a # NOTE: cold restart by design with a brief rationale so the next reader doesn't wonder whether this is a gap.

[thomas-schweich]

Fixed in 92d618b. / helpers added in pawn.trainer; load_resume_state now restores the saved Optax PyTree via unflatten_opt_state(template, flat) when optimizer.safetensors is present and falls back to a fresh template with a loud stderr warning when it isn't. Two new tests in test_jax_lifecycle.py cover both branches.

[thomas-schweich]

Reply text below in body — fix for bash expansion in earlier post.

Fixed in 92d618b. flatten_opt_state / unflatten_opt_state helpers added in pawn.trainer; load_resume_state now restores the saved Optax PyTree via unflatten_opt_state(template, flat) when optimizer.safetensors is present and falls back to a fresh template with a loud stderr warning when it isn't. Two new tests in test_jax_lifecycle.py cover both branches.

[claude]

Optimizer state is never saved, so load_resume_state always cold-restarts Adam. The checkpoint format already has the optimizer.safetensors slot — it just isn't wired up.

Minimal fix:

import equinox as eqx
from pawn.trainer import _flatten_opt_state  # helper you'd add to trainer.py

def _save_checkpoint(step_int: int) -> None:
    out = logger.run_dir / f"step_{step_int:08d}"
    if out.exists():
        return
    opt_tensors = _flatten_opt_state(state.opt_state)   # numpy dict
    save_model(
        state.model, out,
        run_config=cfg.model_dump(),
        optimizer_state=opt_tensors,
        training_state={"step": int(state.step)},
    )
    if push_tracker:
        push_checkpoint_async(out, push_tracker)

Even if cold restart is chosen intentionally, a comment here would prevent a future contributor from wiring it up incorrectly.

[thomas-schweich]

Fixed in 92d618b. _save_checkpoint now passes optimizer_state=flatten_opt_state(state.opt_state) to save_model, so optimizer.safetensors is written into every checkpoint dir alongside model.safetensors and picked up automatically by load_resume_state.

[claude]

Cache key is a hash of the source identifier string, not the model file content. The module docstring says "Cached by source-content hash" and the function is named _content_hash, both implying the actual weights are hashed. In practice, re-publishing thomas-schweich/pawn-base with new weights will not invalidate the cache — the same identifier string produces the same SHA-256, and the old converted checkpoint is returned.

Options:

1. Hash the source file content (pull model.safetensors and hash it) — expensive for large files but truly content-addressed.
2. Hash a cheap proxy (e.g., the HF commit SHA via HfApi().model_info(repo_id).sha) — cheap and accurate.
3. Keep current behaviour but rename to _source_id_hash and update the docstring to say "key is the source identifier, not its content; use force=True if the upstream checkpoint has been updated."

Option 3 is the smallest diff if the semantics are intentional.

[thomas-schweich]

Fixed in 92d618b with option 3 (smallest diff). Renamed _content_hash → _source_id_hash and updated the module docstring + function docstring to spell out that the key is the source identifier (not the content), and that force=True is the documented escape hatch when the upstream HF repo has been re-published. The cheap-proxy idea (HF commit SHA via HfApi().model_info(repo_id).sha) is a reasonable follow-up if we hit real-world cache-staleness — would prefer not to add a new network round-trip to every conversion call right now.

[claude]

WSD boundary can go negative when warmup_steps is explicitly overridden.

_check_lr_schedule_fractions validates warmup_frac + decay_frac ≤ 1.0 in fractions, but _warmup_steps may resolve the warmup to an explicit warmup_steps count that ignores warmup_frac. Example:

total_steps=1000, warmup_steps=600, decay_frac=0.5
→ warmup=600, decay_steps=500
→ stable boundary = 1000 - 500 = 500   # BEFORE the warmup ends at 600

The resulting join_schedules gets non-monotonic boundary points [600, 500], which is undefined behaviour in Optax (schedules are piecewise linear; a negative-width phase will likely clamp to the wrong constant).

Fix: validate the actual resolved step counts in make_lr_schedule, not just the fractions:

if cfg.lr_schedule == "wsd" and warmup + decay_steps > total_steps:
    raise ValueError(
        f"wsd warmup ({warmup}) + decay_steps ({decay_steps}) "
        f"exceeds total_steps ({total_steps})"
    )

Same check is needed for infinite (warmup + cooldown_steps + decay_steps).

[thomas-schweich]

Fixed in 92d618b. make_lr_schedule now raises ValueError for both wsd and infinite schedules when cfg.warmup_steps is not None AND the resolved warmup + decay (or warmup + cooldown + decay) exceeds total_steps. Pure float-fraction rounding overflow (no explicit warmup_steps) is still tolerated by the existing max(0, …) clamp — that's the behaviour the test_lr_schedule_infinite_with_extreme_rounding_doesnt_crash test was asserting. Two new tests cover the rejection path.

[claude]

Zero-length games produce an incorrect loss mask when prepend_outcome=False.

For a game with capped_lengths[i] = 0 (no moves fit in the sequence):

threshold = (0 - 1).clip(min=0) = 0

loss_mask[:, 0] is True (0 <= 0), but tokens[:, 0] is PAD_TOKEN and targets[:, 0] is also PAD_TOKEN (shifted from an all-PAD row). The model is trained to predict PAD from PAD at position 0.

The engine won't produce zero-length games, but pack_corpus is a public API that external callers can feed with any game_lengths. A safer threshold would exclude the empty-game case:

threshold = jnp.maximum(capped_lengths - 1, -1)[:, None]
# -1 ensures seq_positions (≥ 0) > threshold when capped_lengths == 0
loss_mask = seq_positions <= threshold

This is a latent edge case but worth hardening since pack_corpus is documented as a general packing helper.

[thomas-schweich]

Fixed in 92d618b. Threshold for prepend_outcome=False changed to np.maximum(capped_lengths - 1, -1) so zero-length rows get an empty loss mask (positions ≥ 0 are always > -1). New test test_pack_corpus_zero_length_game_has_empty_loss_mask asserts that a row with game_length=0 has loss_mask.any() == False while neighbouring non-empty rows still supervise their real moves.

[claude]

total_steps=None is caught at runtime rather than at config parse time.

PretrainConfig inherits total_steps: int | None = None from BaseRunConfig, so an invocation that omits --total-steps passes pydantic validation, fails silently through all config logic, and only surfaces as a print + return 2 here. Every other mandatory-but-unset issue in the configs raises a ValueError from a @model_validator.

Suggestion — add to PretrainConfig:

@model_validator(mode="after")
def _require_total_steps(self) -> "PretrainConfig":
    if self.total_steps is None:
        raise ValueError(
            "PretrainConfig requires total_steps; pass --total-steps N "
            "or set it in the JSON config"
        )
    return self

This also removes the need for the redundant runtime check in the script.

[thomas-schweich]

Fixed in 92d618b. Added a @model_validator(mode="after") to PretrainConfig._check_pretrain that raises ValueError("PretrainConfig requires total_steps; ...") when the field is None. Removed the now-redundant print + return 2 runtime check in scripts/train_jax.py and updated the test helpers / dispatch tests to supply total_steps=100 on the pretrain happy paths.

[claude]

shutdown(wait=False) can silently abandon in-flight uploads.

HFPushTracker.join calls fut.cancel() on timed-out futures. If a future is already running, Future.cancel() returns False and the thread keeps going. The code counts this as a failure and moves on, then shutdown(wait=False) is called — which also does not wait. The result is a running upload thread (potentially mid-write to HF) abandoned with no cleanup.

For the SIGTERM path this is usually benign (the upload either completes before the process is killed by the kernel, or it doesn't). But under a test that mocks the executor, or on platforms that kill immediately on exit, the partial upload can corrupt the remote branch.

Consider shutdown(wait=True, cancel_futures=False) so already-running uploads always complete, while queued-but-not-started ones are dropped.

[thomas-schweich]

Fixed in 92d618b. HFPushTracker.shutdown now calls self._executor.shutdown(wait=True, cancel_futures=True) — already-running uploads finish (no abandoned thread mid-write to HF), and queued-but-not-started uploads drop so we don't block on a backlog. Docstring updated to spell out the reasoning.

[claude]

Code Review — JAX/Equinox Migration (PR #115)

This is a complete framework swap of the PAWN training stack from PyTorch/Equinox-via-torch to JAX/Equinox/Optax. The architecture is sound, the checkpoint format is well-designed, and the 619-test suite gives reasonable confidence in correctness. Seven inline comments cover specific issues; this summary groups them by severity.

---

🔴 Correctness issues

1. Optimizer state is never saved or loaded on resume (lifecycle.py:263, scripts/train_jax.py:192-202)

_save_checkpoint never passes optimizer_state to save_model, and load_resume_state always calls optimizer.init(...) from scratch. Adam's first/second moment estimates are silently discarded on every resume. For long pretraining runs resumed mid-cosine schedule, the cold restart causes a multi-hundred-step loss spike. The v1 checkpoint format saved optimizer state; the v2 format has the optimizer.safetensors slot but neither end wires it up.

2. WSD/infinite LR schedules can produce non-monotonic boundaries when warmup_steps is explicitly overridden (trainer.py:281-288)

_check_lr_schedule_fractions validates fractions, but _warmup_steps() can return a resolved integer that ignores warmup_frac. Example: total_steps=1000, warmup_steps=600, decay_frac=0.5 → boundary total_steps - decay_steps = 500 falls before the warmup ends at 600. Optax join_schedules with non-monotonic transition points is undefined behaviour. The same gap exists for the infinite schedule (warmup + cooldown + decay can exceed total in step counts even when fractions pass). The fix is to validate the resolved step counts inside make_lr_schedule.

---

🟡 Design/correctness concerns

3. Legacy converter cache key is a source-identifier hash, not a content hash (legacy.py:60-64)

_content_hash(source) hashes the string "1|thomas-schweich/pawn-base". If the upstream checkpoint is republished with new weights, the same key is returned and the stale conversion is used silently. The module docstring and function name both say "source-content hash" which is misleading. Either hash a cheap proxy (HF commit SHA) or rename/document clearly that force=True is required when upstream changes.

4. Zero-length game edge case in pack_corpus (corpus.py:200-207)

For prepend_outcome=False with capped_lengths[i] = 0, threshold = (0-1).clip(min=0) = 0, so loss_mask[:,0] is True even though tokens[:,0] is PAD. The model is trained to predict PAD→PAD. The engine never produces zero-length games, but pack_corpus is a public API. Fix: use max(capped_lengths - 1, -1) as the threshold so zero-length games get an empty loss mask.

---

🟠 Validation / UX

5. total_steps=None for PretrainConfig is caught at runtime, not parse time (scripts/train_jax.py:138-141)

Every other required-but-missing field raises a pydantic ValueError. This one is caught by a print + return 2 in the script. A @model_validator in PretrainConfig would make the error message consistent and machine-readable.

---

🟢 Minor

6. HFPushTracker.shutdown(wait=False) can abandon in-flight uploads (lifecycle.py:109)

join calls fut.cancel() on timed-out futures; if the upload thread is already running, cancel() is a no-op and shutdown(wait=False) then abandons the live thread. Consider shutdown(wait=True, cancel_futures=False) so in-flight uploads always finish, while queued ones are dropped.

---

What looks good

• Checkpoint atomicity (_sentinel.py, checkpoint.py): the .tmp → rename → .complete pattern is solid; the SHA-256 manifest covers all required payloads and the defense against manifest-omission is explicit.
• Supernet slicing (model.py:sliced): inner [:dv, :dv] block semantics are applied consistently across all 16 fields including the non-square FFN weights and the lm_head input axis.
• RoPE in fp32 then downcast, RMSNorm in fp32 then downcast — both match the plan's v2 cast order and are consistent with the parity test contract.
• lax.scan over the layer axis — stacking weights with a leading n_layers axis is the right call at seq=512 where compile time matters.
• LoRA init and scaling (lora.py) — kaiming-uniform A, zero B, alpha/rank scaling, and correct (n_layers, in, out) einsum conventions throughout.
• extra="forbid" on all pydantic models — stale flags surface immediately at config load.
• SIGTERM handler (lifecycle.py:install_sigterm_handler) — idempotent, non-blocking, and correctly decoupled from the training loop.

The DEFERRALS.md "No deferrals" claim is credible: the acceptance-criteria surface (619 passing tests + final_smoke.md live run receipts) is unusually well-documented.