Refactor train.py into giant/training/ around a metrics collector

Every metric name used to exist in four places: the dict keys each
StageTrainer returned, the hardcoded _metrics_fields() column list, the
~110-line metrics_row assembly in train(), and the tqdm/summary
formatting. The two had to be kept in exact correspondence by hand or
csv.DictWriter would raise.

Each metric is now declared once, as a MetricSpec on the trainer that
computes it. MetricsCollector derives the CSV header and W&B payload from
those declarations and owns all accumulation, so train() no longer carries
a running sum, and every isinstance(tr, WGANStageTrainer) branch is gone —
replaced by four trainer hooks (batch_loss, summary, val_objective,
supports_val_loss).

giant/train.py (1875 lines) becomes giant/training/:
  trainers.py       StageSpec + shared StageTrainer base + the two subclasses
  metrics.py        MetricSpec, MetricsCollector
  stage2_inputs.py  the pure AR/teacher-forcing tensor helpers, moved verbatim
  loop.py           train() (225 lines, was ~514) + graceful shutdown
  checkpoint.py     build/load, lifted out of train()'s closures

The trainers shared ~15 identical constructor arguments and copy-pasted
their cosine-warmup lambda, EMA setup, state_dict/load_state_dict,
resume_lr and train_mode/eval_mode. StageSpec resolves one stage's config
once (constructors go from 24 and 22 keyword arguments to (spec, model,
device)), the base class holds the rest, and build_stage_trainers drops
from ~100 lines to 15.

Metric columns are renamed to a uniform stage/split/metric scheme
(stage1/train/loss, stage2/train/d_loss, stage1/lr, stage1/router/entropy,
val/loss, ...). Old metrics.csv files and W&B history are not comparable.
The checkpoint format is unchanged.

BEHAVIOR CHANGE — WGAN best-checkpoint selection. The old code meant to
score a WGAN stage on its marginal KL, but the guard
`{n: kl for n in wgan_names if n not in val_loss_per_stage}` could never
fire: val_loss_per_stage was pre-seeded with 0.0 for every stage, so a
WGAN stage contributed a flat 0.0 and the KL was written to metrics.csv
without ever influencing best.pt. val_objective now returns it as
intended. On the test harness's default flow+wgan config val_loss went
from 2.182 (stage 1 only) to 15.137 (stage 1 + KL 12.954), and which epoch
won changed. Runs before this commit picked their best checkpoint on the
non-adversarial stages alone. Written up in docs/v0.3.0-followups.md.

Verified: 699 tests pass; ruff, ruff format and ty clean. Baseline-vs-
refactor metrics.csv compared across five configs (flow+wgan, AR+onehot,
routed, both-flow, AR-flow) — every comparable value bit-identical except
val/loss where the fix applies. Resume appends without a duplicate header
and reproduces a HEAD worktree's per-epoch losses and LRs exactly across
the resume boundary. A refactored last.pt loads through
cli.py:_load_model_weights in both raw and ema modes.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

docs/v0.3.0-design.md

@@ -322,7 +322,7 @@ stage1_context = "truth"
 | `lambda` | float | `1.0` | Weight of stage 2's loss in the total. Was `train.lambda_s2`. |
 | `k_max` | int | `15` | Maximum secondary slots. Under `"one_shot"` this is the fixed output width; under `"autoregressive"` it is a safety cap on the generation loop. Was the global constant `K_MAX` in `giant/constants.py` (max observed `n_sec` is 14 in the PbWO4 dataset, so 15 covers it with one spare). |
 | `context_dim` | int | `64` | Width of the projected stage-1 outcome fed into stage 2's conditioning. Was the hardcoded `stage1_proj_dim = 64`. |
-| `stage1_context` | `"truth"` \| `"sampled"` | `"truth"` | What stage 2 conditions on during training. `"truth"`: the ground-truth stage-1 target vector, detached — v0.2 behaviour (`train.py:254` passes `x1_s1.detach()`), i.e. stage-level teacher forcing. `"sampled"`: stage 1's own sampled output, closing the train/inference gap at the cost of a sampling pass per batch and a moving target early in training. |
+| `stage1_context` | `"truth"` \| `"sampled"` | `"truth"` | What stage 2 conditions on during training. `"truth"`: the ground-truth stage-1 target vector, detached — v0.2 behaviour (v0.2 `train.py:254` passes `x1_s1.detach()`), i.e. stage-level teacher forcing. `"sampled"`: stage 1's own sampled output, closing the train/inference gap at the cost of a sampling pass per batch and a moving target early in training. |
 
 #### `[stage2_model.n_sec]`
 
@@ -748,7 +748,7 @@ configuration that plausibly meets the budget; flow AR is for quality comparison
 
 ## 7. Training loop
 
-Decision 2 (full mixed per-stage objectives) makes `train.py` one trainer object
+Decision 2 (full mixed per-stage objectives) makes `giant/training/` one trainer object
 per active stage:
 
 ```python
@@ -765,8 +765,15 @@ FlowTrainer, DDPMTrainer, WGANTrainer(critic, n_critic, gp_weight)
   critic updates then a generator update. Independent optimizers, independent EMA.
 - Router auxiliary losses (`lambda_balance` / `lambda_proc` / `lambda_entropy`)
   become per-stage, summed over whichever stages are routed. Today's
-  `hasattr(model, "router")` check (`train.py:262`) generalizes cleanly.
-- `metrics.csv` and W&B metric names gain a stage prefix.
+  `hasattr(model, "router")` check (v0.2 `train.py:262`) generalizes cleanly.
+- `metrics.csv` and W&B metric names gain a stage prefix. **Implemented as**
+  `<stage>/<split>/<metric>` (`stage1/train/loss`, `stage2/train/d_loss`,
+  `stage1/lr`, `stage1/router/entropy`) plus an unprefixed run-level tail
+  (`val/loss`, `val/marginal_kl`, `grad_norm`, `gpu_mem_mb`,
+  `samples_per_sec`, `is_best`, `epoch_time_s`). Each name is declared once,
+  as a `MetricSpec` on the `StageTrainer` that computes it; the CSV header
+  and W&B payload are derived from those declarations
+  (`giant/training/metrics.py`).
 - With `stage1_model.active = false`, stage 2 still needs its stage-1 context: it
   comes from the ground-truth target already in the batch (`x1_s1`), which is
   exactly what v0.2 does anyway. **Stage-2-only training is therefore a cheap
@@ -855,7 +862,7 @@ All in `giant/config.py`:
 | File | Why |
 |------|-----|
 | `giant/pipeline.py` | Builds `model_config`; now per-stage. Delete the wgan+router rejection at `:275`. Router `centers_init` seeding becomes per-stage. |
-| `giant/train.py` | Per-stage trainers (§7). |
+| `giant/training/` | Per-stage trainers (§7), metric collection, checkpointing, the epoch loop. |
 | `giant/cli.py` | Stage-prefixed flags for `train` and `new-run`; `build_models` now returns a dict (`:871`, `:1254`); `:1339` writes `model_config` into the rollout sidecar. |
 | `giant/sample.py` | Sampler picked per stage from `stage*_model.generator`; new AR sampling loop with KV cache under `history = "attention"`. |
 | `giant/rollout.py` | AR secondary generation; categorical class -> PDG decode; `other_policy` handling. |
@@ -894,7 +901,7 @@ All in `giant/config.py`:
   later is a config addition, not a break.
 - **`stage2_model.generator = "ddpm"`.** The value is **accepted by the
   schema** (§3.3 lists `"flow" | "ddpm" | "wgan"` with no caveat) but
-  `FlowDDPMStageTrainer.__init__` (`giant/train.py`) raises
+  `FlowDDPMStageTrainer.__init__` (`giant/training/trainers.py`) raises
   `NotImplementedError` for stage 2 — only `"flow"` and `"wgan"` have a
   stage-2 secondary-decoder loss implemented. `stage1_model.generator =
   "ddpm"` is unaffected; this restriction is stage-2-only. Landing stage-2
@@ -968,7 +975,7 @@ or a non-adversarial CE head, both of which already exist as config options.
 2. **`network.py`** — the §5 decomposition, with `Stage2OneShot` reproducing v0.2
    exactly. Gate on the §4.3 migration test: load a v0.2 checkpoint through the
    shim and diff outputs against v0.2 code.
-3. **`train.py`** — per-stage trainers; `active = false` paths. At this point
+3. **`giant/training/`** — per-stage trainers; `active = false` paths. At this point
    Stage-2-only training works and the meeting's step 2 (one-shot WGAN baseline,
    trained standalone) is runnable.
 4. **Type map** — `loader.py` + `setup_cache.py` + checkpoint persistence +

docs/v0.3.0-followups.md

@@ -30,8 +30,24 @@ turned out fine:
 - `analysis/render.py`/`analysis/router_gating.py` correctly branch
   old-flat vs new-nested `model_config["router"]` location — doc flagged this
   as a likely stale spot (§10) but it's actually fine.
-- `train.py`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence,
+- `giant/training/`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence,
   per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only
   training via ground-truth `x1_s1` (§7).
+
+## Behavior change: WGAN best-checkpoint selection
+
+The `giant/training/` split fixed a dead guard in WGAN validation scoring. A
+WGAN stage was *meant* to contribute its marginal KL to the `val_loss` that
+drives `best.pt`, but the guard `if n not in val_loss_per_stage` could never
+fire (every stage was pre-seeded to `0.0`), so the stage contributed a flat
+`0.0` and the KL was written to `metrics.csv` without ever being used.
+`WGANStageTrainer.val_objective` now returns the KL as intended.
+
+**Consequence:** any checkpoint selected before this commit under a config
+with a WGAN stage — including the v0.3.0 default (`stage2_model.generator =
+"wgan"`) — picked its best epoch on the non-adversarial stages alone. Measured
+on the test harness's default flow+wgan config, `val_loss` went from `2.182`
+(stage 1 only) to `15.137` (stage 1 + KL `12.954`), and which epoch won
+changed. Do not compare `val/loss` or `best.pt` choice across this commit.
 - `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init`
   seeding, removed stale expert-size warning (§9, §10).