v0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite

Replaces train.py's single global training loop with a StageTrainer
hierarchy (FlowDDPMStageTrainer, WGANStageTrainer) — one per active
stage, each owning its own optimizer/LR schedule/EMA and reading only
the shared batch tuple (stage 2 always teacher-forces on the
ground-truth x1_s1, so stages never need each other's output at train
time). Supports every stage1/stage2 generator combination, including
the design doc's headline mixed case (stage1=flow + stage2=wgan) and
its reverse, plus stage1-only/stage2-only ablation runs, routed+gumbel
stages, and checkpoint save/resume. metrics.csv/wandb logging are
stage-prefixed. validate_marginals calls are guarded with a one-time
warning and a Wasserstein-magnitude fallback for wgan best-checkpoint
selection, since giant/sample.py still assumes stage1 always owns
n_sec_head (decision 1 moved it to stage 2 by default) — deferred to
design doc step 6, not silently papered over.

pipeline.py's run_setup_stage/run_train_job now read the new nested
config directly; the dangling resolve_expert_dims call and the
--mode wgan --router rejection are both gone (routed WGAN works).
cli.py's train/new-run build correctly-shaped config overrides
(architecture flags -> stage1_model only per the approved decision;
--mode/--n-critic/--gp-weight/--critic-lr broadcast to both stages,
matching migrate_config's own precedent and avoiding a regression on
the common --mode case); predict/rollout's dangling build_models
tuple-unpack is fixed; new-run now tags config_version, fixing a bug
where a re-loaded v0.3 config.toml would have been silently corrupted
by migrate_config mistaking it for v0.2.

config.py's validate_config rejects mixed particle/material
conditioning types for now (ConditionEncoder supports it, the data
pipeline in giant/data/transforms.py doesn't yet). analysis/render.py
and router_gating.py handle both the new nested model_config shape and
legacy flat checkpoints. scripts/warm_setup_cache.py updated for
run_setup_stage's new signature.

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

scripts/warm_setup_cache.py

@@ -39,12 +39,23 @@ def run_warm_setup_cache(
         "type": router_type,
         "n_experts": n_experts,
     }
+    # A minimal v0.3 cfg — only the keys run_setup_stage actually reads
+    # (conditioning.particle.type, stage{1,2}_model.router). This CLI only
+    # ever configures one router (matching today's single --router-type
+    # flag), so it's placed on stage1_model; stage2_model's stays disabled.
+    cfg = {
+        "conditioning": {
+            "particle": {"type": conditioning},
+            "material": {"type": conditioning},
+        },
+        "stage1_model": {"router": router_cfg},
+        "stage2_model": {"router": {"enabled": False}},
+    }
     run_setup_stage(
         Path(data),
         val_fraction=val_fraction,
         seed=seed,
-        conditioning=conditioning,
-        router_cfg=router_cfg,
+        cfg=cfg,
         cache_setup=True,
         rebuild_setup_cache=rebuild,
         echo=echo,