Exact one-epoch training from local data (padding-free distributed sampler)

[ksd3]

For scaling studies that must train every config on the identical data, every example exactly once. Streaming + split_dataset_by_node + buffered shuffle can't do that — the order depends on GPU/worker count, and the standard distributed samplers either drop the tail or pad it with duplicates when len % world_size != 0.

src/astropt/exact_epoch.py (new, unit-tested)

• ExactDistributedSampler — partitions range(N) across ranks with no padding, no dropping: a single seeded shuffle, sliced strided (idx[rank::world_size]). Exact coverage, reproducible (depends only on data_seed), and independent of num_workers (the sampler fixes the order; workers only fetch).
• one_epoch_loop — consumes the loader exactly once; an optimizer step every grad_accum_per_rank micro-batches with a short final step. The exact partition gives every rank the same micro-batch count, so the short final step is identical across ranks → the per-step all-reduce count stays balanced and DDP needs no Join. Gradients sync every backward (all-reduce is linear ⇒ identical to accumulate-then-sync).
• Self-tests (python -m astropt.exact_epoch): exact partition for many sizes incl. a prime; the loop consumes exactly N with the expected step count. For the galaxies set (N=8,474,566), W∈{1,2,4,8} all balance to 13,242 steps.

scripts/train.py

When data_dir is set: load the galaxies map-style from local parquet, shard with ExactDistributedSampler, run one_epoch_loop, and assert the all-reduced example count == dataset size — a run self-proves exact coverage or fails loudly. Streaming path unchanged.

Notes

• Adds data_seed (also introduced by the data-order PR Reproducible streaming data order across runs (fixed data_seed + pinned num_workers) #70 — trivial overlap to resolve on merge; the streaming-shuffle block is also restructured into if exact / elif use_hf / else).
• The DDP path is arithmetically balanced and self-asserting, but I could not run it multi-GPU — it needs a smoke test on a real 4-GPU node before relying on it.