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>
Pulled forward from the not-yet-mergeable condor-gpu-train-rollout branch:
`new-run` resolves CLI hyperparameter overrides into a full config.toml and
run dir (reusing the existing default_out_dir_name collision-avoidance and a
newly factored-out router-override helper shared with `train`), so a run can
be prepared and reviewed before `giant train` actually kicks off. Also
brings README up to date with the model/CLI as it actually stands
(physical/embedding conditioning, WGAN/MoE-router modes, giant analyze,
W&B, setup-stage caching), which had drifted back to describing the
Phase-1 proof-of-concept.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
lars