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>
Replace the single global train.mode + [model] block with the four
top-level blocks docs/v0.3.0-design.md specifies ([conditioning],
[stage1_model], [stage2_model], [train]), so Stage 1 and Stage 2 can
run independent generative objectives and Stage 2 can train standalone.
- migrate_config translates old config.toml/checkpoint dicts on load,
so nothing on /ceph goes dead; loudly rejects non-zero
expert_hidden_dim/expert_n_blocks, which v0.3.0 no longer supports.
- merge_cli_overrides/save_config generalize from one hardcoded nesting
level (model.router) to arbitrary recursive depth.
- default_out_dir_name candidates move to dotted paths against the new
schema, with per-stage router/generator discriminators.
- validate_config adds cross-block checks the per-block schema can't
express (particle_type.target=embedding needs a matching conditioning
mode, tie_to_stage1 needs an active stage 1, etc).
- resolve_expert_dims is deleted (experts always inherit the stage's
hidden_dim/n_res_blocks now) — pipeline.py/cli.py callers are left
dangling on purpose, to be updated in the network.py/train.py steps
that follow.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Extends default_out_dir_name's non-default-field convention to the
router's new gumbel combine-weight flag and its learnable-knob toggles,
so gumbel sweep configs (learn_centers on/off, learn_width,
learn_temperature) resolve to distinguishable checkpoint directory
names instead of colliding on the same r-<type><n> token.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Trains the routed trunk's forward combination as a hard one-hot sample
(matching eval-time top-1 dispatch exactly) while keeping a smooth gradient
on the backward pass, targeting the train/eval mismatch identified as a
likely contributor to experts overlapping instead of partitioning in the
first energy-router rollout benchmark. Off by default (model.router.gumbel);
existing routed configs/checkpoints are unaffected.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Previously the same fixed 7 fields (mode/hidden_dim/n_blocks/emb_dim/
conditioning/lr/batch_size) were always baked into the name, even for a
vanilla run, and router config wasn't represented at all. Now
default_out_dir_name only includes fields that differ from
DEFAULT_CONFIG, adds router/seed/epochs as candidates, and caps at 6
shown fields with a hashed overflow suffix for heavily-swept configs.
Covers the default-config 0-sentinel inheritance path (the exact bug
fixed by 969c5c6, previously untested since every router test in
test_router.py passes expert_hidden_dim/expert_n_blocks explicitly),
plus missing-key inheritance, full override, and partial override.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Provides stratified marginal comparisons, joint-structure checks (correlation
matrices, physically-coupled pairwise plots, direction alignment), and
physical-constraint validation (unit-norm directions, non-negative raw
targets) for a trained model's generated samples, building on the aggregate
marginal/KL check already in giant.validate.
Supports two entry points: live sampling against a checkpoint + val data
(load_model_bundle/collect_samples), or loading a precomputed
`giant predict --coord local` parquet directly (load_predicted_local) without
needing the checkpoint at all. Predict output is now tagged with parquet
schema metadata so the loader can verify a file's format and reject
coord=global or untagged files with a clear error instead of guessing from
column names.
Also extends the config git-hash mismatch warning (added for --config
loading) to checkpoint loading: both `giant predict` and
analysis.load_model_bundle now look for a config.toml next to the checkpoint
and warn (without failing) if it was generated from a different git commit.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
lars