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>

tests/test_train.py

@@ -1,4 +1,4 @@
-"""Tests for giant/train.py."""
+"""Tests for giant/training/."""
 
 import copy
 import csv
@@ -18,14 +18,19 @@ from giant.constants import (
     X_DIM,
 )
 from giant.model.network import build_critics, build_models
-from giant.train import (
+from giant.training import (
     FlowDDPMStageTrainer,
+    StageSpec,
     WGANStageTrainer,
+    build_stage_trainers,
+    train,
+)
+from giant.training.metrics import _wandb_run_config
+from giant.training.stage2_inputs import (
     _ar_has_prev,
     _assemble_stage2_ar_inputs,
     _assemble_stage2_ar_target,
     _assemble_stage2_real,
-    _build_stage_trainers,
     _gumbel_tau,
     _relax_onehot_type_slice,
     _remaining_energy_fraction,
@@ -33,8 +38,6 @@ from giant.train import (
     _stage2_tf_prob,
     _stick_fraction,
     _type_repr,
-    _wandb_run_config,
-    train,
 )
 
 PDG_VOCAB = 6
@@ -558,9 +561,9 @@ def test_metrics_csv_columns_are_stage_prefixed():
         out_dir = Path(tmp) / "run"
         _run_train(cfg, out_dir)
         header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
-        assert "stage1_train_loss" in header
-        assert "stage2_train_d_loss" in header
-        assert "val_loss" in header
+        assert "stage1/train/loss" in header
+        assert "stage2/train/d_loss" in header
+        assert "val/loss" in header
         assert "epoch" in header
 
 
@@ -574,22 +577,16 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
     models = build_models(model_config)
     critics = build_critics(model_config)
     assert models["stage1"] is not None and critics["stage1"] is not None
-    trainer = WGANStageTrainer(
+    spec = StageSpec(
         name="stage1",
-        model=models["stage1"],
-        critic=critics["stage1"],
         is_stage2=False,
-        lambda_weight=1.0,
-        n_sec_lambda=0.1,
+        generator="wgan",
         n_critic=1000,  # never a generator step in this test
-        gp_weight=10.0,
-        lr=3e-4,
-        critic_lr=0.0,
         ema_decay=0.0,
-        warmup_epochs=0,
-        epochs=1,
         steps_per_epoch=4,
-        device=torch.device("cpu"),
+    )
+    trainer = WGANStageTrainer(
+        spec, models["stage1"], critics["stage1"], torch.device("cpu")
     )
     assert trainer.model.n_sec_head is None
     batch = _fake_batches(1, 8)[0]
@@ -598,28 +595,11 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
 
 
 def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
+    spec = StageSpec(
+        name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0
+    )
     with pytest.raises(NotImplementedError):
-        FlowDDPMStageTrainer(
-            name="stage2",
-            model=torch.nn.Linear(1, 1),
-            is_stage2=True,
-            generator="ddpm",
-            lambda_weight=1.0,
-            n_sec_lambda=0.1,
-            lambda_balance=0.0,
-            lambda_proc=0.0,
-            lambda_entropy=0.0,
-            gumbel_tau_start=1.0,
-            gumbel_tau_end=0.1,
-            lr=3e-4,
-            weight_decay=0.01,
-            ema_decay=0.0,
-            warmup_epochs=0,
-            epochs=1,
-            steps_per_epoch=1,
-            ddpm_n_steps=50,
-            device=torch.device("cpu"),
-        )
+        FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
 
 
 # --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
@@ -649,7 +629,7 @@ def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
     model_config = _model_config(cfg)
     models = build_models(model_config)
     critics = build_critics(model_config)
-    trainers = _build_stage_trainers(
+    trainers = build_stage_trainers(
         cfg, models, critics, torch.device("cpu"), total_train_batches=4
     )
     trainer = trainers["stage2"]
@@ -686,7 +666,7 @@ def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
             rows = list(csv.DictReader(f))
         assert len(rows) == cfg["train"]["epochs"]
         loss_col = (
-            "stage2_train_g_loss" if stage2_generator == "wgan" else "stage2_train_loss"
+            "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
         )
         assert all(math.isfinite(float(r[loss_col])) for r in rows)
 
@@ -705,10 +685,10 @@ def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics():
         _run_train(cfg, out_dir)
         with open(out_dir / "metrics.csv", newline="") as f:
             rows = list(csv.DictReader(f))
-        assert "stage2_train_grad_norm_type_slice" in rows[0]
-        assert "stage2_train_grad_norm_cont_slice" in rows[0]
-        assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows)
-        assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows)
+        assert "stage2/train/grad_norm_type_slice" in rows[0]
+        assert "stage2/train/grad_norm_cont_slice" in rows[0]
+        assert any(float(r["stage2/train/grad_norm_type_slice"]) > 0 for r in rows)
+        assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows)
 
 
 def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
@@ -721,8 +701,8 @@ def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
         _run_train(cfg, out_dir)
         with open(out_dir / "metrics.csv", newline="") as f:
             rows = list(csv.DictReader(f))
-        assert any(float(r["stage2_train_grad_norm_type_slice"]) > 0 for r in rows)
-        assert any(float(r["stage2_train_grad_norm_cont_slice"]) > 0 for r in rows)
+        assert any(float(r["stage2/train/grad_norm_type_slice"]) > 0 for r in rows)
+        assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows)
 
 
 def test_wgan_physical_omits_grad_norm_slice_columns():
@@ -731,5 +711,5 @@ def test_wgan_physical_omits_grad_norm_slice_columns():
         out_dir = Path(tmp) / "run"
         _run_train(cfg, out_dir)
         header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
-        assert "stage2_train_grad_norm_type_slice" not in header
-        assert "stage2_train_grad_norm_cont_slice" not in header
+        assert "stage2/train/grad_norm_type_slice" not in header
+        assert "stage2/train/grad_norm_cont_slice" not in header