fix(harvest): per-worker local_batch_size; GPU assert + mem-fraction for 8B restores

[ocg-goodfire]

What

Three fixes from the first attempt to harvest an 8B JAX run (p-289e3178, L18-MLP 200k) on the sharded SLURM path:

1. HarvestConfig.batch_size → local_batch_size (per-worker). The sharded launcher divided the global batch_size across the array (ShardServer process_count=n_gpus), so any production config quietly ran every GPU at a local batch of 1 — e.g. batch_size: 128 on 128 GPUs. Workers now always forward local_batch_size sequences per batch and the schedule's global batch derives as local_batch_size × world_size; total sequences = n_batches × local_batch_size × world_size. No back-compat alias — old configs fail loudly on the removed field.

2. Fail-fast GPU assert for sharded workers. A venv without CUDA jaxlib fell back to CPU silently; a 64-GPU array burned its allocation at ~100× slowdown with nothing in the logs but a jax warning. Sharded workers (--rank given) now assert jax.default_backend() == "gpu". Single-process CPU spot-checks stay allowed.

3. XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 default in the worker. open_jax_run restores the full TrainState (fp32 masters + Adam + persistent-PGD sources — 60GB checkpoint for the L18 runs) plus the frozen 8B target onto one device; restore staging pushes the peak past the default 0.75 pool (RESOURCE_EXHAUSTED at ~135GB on a B200). Set before backend init, setdefault so callers can override.

A structural follow-up worth considering separately: open_jax_run could partially restore only components + ci_fn (orbax 0.12 partial_restore) — the optimizer states and adversary sources it materializes are dead weight for every consumer, and dropping them would halve harvest GPU pressure. Kept out of this PR to stay minimal.

Testing

• basedpyright param_decomp_lab/harvest/ clean; harvest test suite (62) passes.
• 1-GPU probe of the exact worker command on p-289e3178 confirmed the OOM and the 0.95 fix; overnight 128-GPU harvests of p-289e3178/p-fc8b6c7a run from this branch (results by morning).

🤖 Generated with Claude Code

https://claude.ai/code/session_01PVGsKRCYkqeeVhizuPnVuD