feat(training): single-node data parallelism via leech model train --gpus N

[jayhesselberth]

leech model train --gpus N runs N data-parallel ranks on one node. Opt-in,
single-node, train only — eval test is input-bound and gets nothing from it.
Closes #241.

The two decisions worth reviewing

--batch-size is the GLOBAL batch, split across ranks. This is the opposite
of the PyTorch convention, and it is what keeps a multi-GPU run comparable with
the arms already measured: the optimizer-step count, the LR schedule,
ClipGrad's quantile buffer and the grad_accum_split arithmetic are all
unchanged, so the same command line means the same recipe at any --gpus.

The weighted sampler shards one global draw rather than drawing per rank.
Every rank draws the same multinomial from a private generator seeded
seed + epoch, then takes draw[rank::world_size]. The union of the shards is
exactly the single-GPU epoch, so the class ratio --oversample-minority
produces is reproduced by construction rather than in expectation.

Measured scaling

Production charging corpus (6.72M chunks, 36.5 GiB npz), 2 epochs, one 4×A30
node. Full numbers and caveats in #241.

--gpus	job wall	speedup	peak RSS
1	24:31	1.00x	40.6 GiB
2	16:40	1.47x	88.3 GiB
4	10:53	2.25x	157.6 GiB

2.25x rather than 4x: ~3 min of per-rank corpus load and the final eval do not
shard, so they floor a 2-epoch run. Training time alone goes 18:30 → 10:14 →
8:55. Memory is ~1:1 with the corpus per rank, so --gpus and mem_mb move
together or the second rank is OOM-killed during its load.

Three corrections to the issue

1. The LR schedule here is epoch-indexed, not step-indexed. Nothing in the
   trainer counts optimizer steps, so criterion 4's second half needed stating
   rather than fixing.
2. The auxiliary heads are a fourth silent failure the issue does not list.
   AdversarialHead and RegressionHead are in the optimizer but outside the
   model, so a DDP around the model alone never allreduces their gradients.
   Each now gets its own wrapper, with a test that fails if one is added without.
3. Weighted CrossEntropyLoss does not decompose over shards. It normalizes
   by summed weight, not count, so each rank divided by its own shard's class
   mix — 8% off on a fixture whose shards differ, silently. Ranks now form
   world_size × local / global weight sum.

Two defects only the hardware found

Spawned ranks were silent. setup_logging runs in the click entry point,
which a rank never reaches, so every INFO line was dropped — including rank 0's.

DataLoader workers were spawning, not forking. multiprocessing.spawn.prepare
forces a spawned child's start method to match how it was created, so every
DataLoader pickled the dataset's stacked tensors through /dev/shm instead of
COW-sharing them — voiding the invariant those contiguous buffers exist for. It
does not look like memory: shm is charged to the cgroup but not to RSS, so the
run died at 176 GiB RSS against a 244 GiB allocation, in the fourth rank's
validation loader, as No space left on device. Fixing it took the same run to
157.6 GiB.

Also here, by request

escapepod v0.16.1 → v0.21.0, on both backends. The crate and the PyPI
package are one upstream released in lockstep, and leech drives both, so they
have to move together — dependabot's #236 bumped the crate alone and left main
skewed. Validated by tests/test_backend_parity.py (86 tests) against a
leech_core actually built on v0.21.0, not just cargo check.

A daily check so that cannot recur (.github/workflows/escapepod-sync.yml):
compares the two pins, and when they disagree rewrites both and opens a PR.
Deliberately daily-and-propose rather than red-on-dependabot's-PR, which would
leave a human to do the second half by hand. uv.lock's entry is rewritten from
PyPI's release record because uv lock — including --upgrade-package —
re-serializes the whole file at ~670 changed lines and buries the real edit.

Acceptance criteria

#	Criterion	Where
1	Numerical equivalence	test_two_rank_gradients_match_the_single_rank_gradients — one step from identical init, 2x4 vs 1x8, equal to 1e-6
2	Sharding verified, not assumed	test_weighted_shards_interleave_to_the_single_global_draw (union == the reference draw), test_weighted_shards_differ_between_ranks
3	Checkpoint byte-compatible	test_two_rank_checkpoint_has_the_single_rank_keys — key-set equality, no module.
4	no_sync() + scheduler	test_no_sync_covers_every_micro_step_but_the_last; scheduler is epoch-indexed (correction 1)
5	Single-GPU unchanged	test_single_gpu_creates_no_process_group; every change is behind world_size > 1
6	Scaling reported honestly	table above, and #241

25 distributed tests run on gloo over CPU, so they run in CI rather than only on
a GPU node. Full suite 1559 passed / 44 skipped.

Not in this PR

The 1 vs 2 GPU same-seed convergence run (~8 GPU-hours) — worth scheduling
deliberately rather than folding into review. The AUROC spread in the table is
not that experiment: augmentation noise is drawn per DataLoader worker, so a
different rank count means a different realized trajectory by construction.

🤖 Generated with Claude Code

https://claude.ai/code/session_01CdbVo6UnMuYNmt2D7tjT3B