From 9112e845e0d168a16d7101f9a827625bc2e03d16 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 6 Aug 2026 11:31:49 +0200 Subject: [PATCH] v0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- giant/analysis/render.py | 52 +- giant/analysis/router_gating.py | 33 +- giant/cli.py | 200 +++- giant/config.py | 12 + giant/pipeline.py | 204 ++-- giant/train.py | 1821 ++++++++++++++++--------------- scripts/warm_setup_cache.py | 15 +- tests/test_cli_new_run.py | 9 +- tests/test_config.py | 18 + tests/test_pipeline.py | 11 +- tests/test_rollout.py | 5 +- tests/test_router.py | 23 +- tests/test_router_gating.py | 3 +- tests/test_train.py | 400 +++++-- 14 files changed, 1688 insertions(+), 1118 deletions(-) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index 3941981..b746942 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -41,11 +41,47 @@ def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> Non ax.set_yscale("log") -def _router_summary(model_config: dict) -> str: - r = model_config.get("router") or {} - if not r.get("enabled"): +def _router_summary(router_cfg: dict) -> str: + if not router_cfg.get("enabled"): return "off" - return f"{r.get('type', '?')}×{r.get('n_experts', '?')}" + return f"{router_cfg.get('type', '?')}×{router_cfg.get('n_experts', '?')}" + + +def _figure_params_v2(mc: dict, run_meta: dict) -> dict: + """`_figure_params` for a new-shape (nested) `model_config` — has a + `stage1_model` key. Reports stage 1's architecture (the headline + generator); stage 2's generator is only added (`mode_s2`) when it + differs from stage 1's, since a mixed run (docs/v0.3.0-design.md's + `stage1=flow` + `stage2=wgan` case) is the interesting exception, not + the common case.""" + s1 = mc["stage1_model"] + s2 = mc.get("stage2_model") or {} + mode = s1.get("generator") + params: dict = {} + if s1.get("hidden_dim") is not None: + params["hidden_dim"] = s1["hidden_dim"] + if s1.get("n_res_blocks") is not None: + params["n_res_blocks"] = s1["n_res_blocks"] + if mode is not None: + params["mode"] = mode + s2_mode = s2.get("generator") + if s2_mode is not None and s2_mode != mode: + params["mode_s2"] = s2_mode + particle_type = ((mc.get("conditioning") or {}).get("particle") or {}).get("type") + if particle_type is not None: + params["conditioning"] = particle_type + params["router"] = _router_summary(s1.get("router") or {}) + if run_meta.get("training_epoch") is not None: + params["epoch"] = run_meta["training_epoch"] + if run_meta.get("best_val_loss") is not None: + params["best_val_loss"] = round(run_meta["best_val_loss"], 4) + if mode == "wgan": + noise_dim = (s1.get("wgan") or {}).get("noise_dim") + if noise_dim is not None: + params["noise_dim"] = noise_dim + elif run_meta.get("steps") is not None: + params["steps"] = run_meta["steps"] + return params def _figure_params(run_meta: dict) -> dict: @@ -59,8 +95,14 @@ def _figure_params(run_meta: dict) -> dict: architecture-conditional: flow/ddpm runs show the ODE ``steps`` used for this rollout, wgan runs show ``noise_dim`` instead since wgan sampling is single-pass and has no ODE step count. + + Handles both a v0.2 checkpoint's flat ``model_config`` and a v0.3.0 + nested one (has a ``stage1_model`` key — see ``_figure_params_v2``). """ mc = run_meta.get("model_config") or {} + if "stage1_model" in mc: + return _figure_params_v2(mc, run_meta) + mode = mc.get("mode") params: dict = {} if mc.get("hidden_dim") is not None: @@ -71,7 +113,7 @@ def _figure_params(run_meta: dict) -> dict: params["mode"] = mode if mc.get("conditioning") is not None: params["conditioning"] = mc["conditioning"] - params["router"] = _router_summary(mc) + params["router"] = _router_summary(mc.get("router") or {}) if run_meta.get("training_epoch") is not None: params["epoch"] = run_meta["training_epoch"] if run_meta.get("best_val_loss") is not None: diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py index 6dbc3f8..9dd6a96 100644 --- a/giant/analysis/router_gating.py +++ b/giant/analysis/router_gating.py @@ -65,6 +65,19 @@ class _RouterHandle: router_type: str +def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str: + """The single conditioning-mode string `giant.data.transforms. + build_cond_features` still expects — from either a v0.2 checkpoint's + flat `model_config["conditioning"]` (already a string) or a new-format + one (`model_config["conditioning"]["particle"]["type"]`; `validate_config` + guarantees particle/material agree, see its mixed-conditioning check). + Mirrors `giant.cli._conditioning_str`.""" + raw = model_cfg.get("conditioning", default) + if isinstance(raw, dict): + return raw.get("particle", {}).get("type", default) + return raw + + def load_router(checkpoint: str | Path) -> _RouterHandle | None: """Load a checkpoint's Stage-1 router, or None if it isn't a MoE checkpoint.""" import torch @@ -74,20 +87,32 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None: ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) model_cfg = ckpt.get("model_config") or {} - router_cfg = model_cfg.get("router") + # New nested shape (has a "stage1_model" key) vs. a v0.2 checkpoint's + # flat model_config. + router_cfg = ( + (model_cfg.get("stage1_model") or {}).get("router") + if "stage1_model" in model_cfg + else model_cfg.get("router") + ) if not router_cfg or not router_cfg.get("enabled"): return None - stage1, _ = build_models(model_cfg) + built = build_models(model_cfg) + stage1 = built["stage1"] + if stage1 is None: + return None stage1.load_state_dict(ckpt["model"]) stage1.eval() + router = stage1.trunk.router + if router is None: + return None return _RouterHandle( - router=stage1.router, + router=router, pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()}, mat_map={str(k): v for k, v in ckpt["mat_map"].items()}, cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]), - conditioning=model_cfg.get("conditioning", "embedding"), + conditioning=_conditioning_str(model_cfg), router_type=router_cfg["type"], ) diff --git a/giant/cli.py b/giant/cli.py index 070a00c..1abf63b 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -74,27 +74,53 @@ def _router_total_experts(router_cfg: dict) -> int: return int(router_cfg.get("n_experts", 1)) -def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int]: +def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str: + """The single conditioning-mode string `giant.data.transforms. + build_cond_features` still expects — from either a v0.2 checkpoint's + flat `model_config["conditioning"]` (already a string) or a new-format + one (`model_config["conditioning"]["particle"]["type"]`; `validate_config` + guarantees particle/material agree, see its mixed-conditioning check).""" + raw = model_cfg.get("conditioning", default) + if isinstance(raw, dict): + return raw.get("particle", {}).get("type", default) + return raw + + +def _batch_size_estimate_dims( + model_cfg: dict, training: bool, stage: str = "stage1" +) -> tuple[int, int]: """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory. - Routed models spend their FLOPs in the (smaller) expert trunks, not the - monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert - dims instead when routing is enabled. Training runs the full soft mixture - (every expert on the whole batch), so its activation memory scales with - the expert count; inference does top-1 dispatch (each row hits one - expert), so the batch just partitions across experts and one expert's - dims already bound it. estimate_batch_size scales memory linearly with - hidden_dim * n_blocks, so the training multiplier folds into n_blocks. + `model_cfg` is either the new nested shape (has a `f"{stage}_model"` key + — the merged training `cfg`, or a checkpoint's new-format `model_config`) + or a v0.2 checkpoint's flat `model_config`. v0.3.0 dropped per-expert + sizing (giant.model.network's routed trunks always inherit the stage's + own hidden_dim/n_res_blocks — no more `resolve_expert_dims`), so the new + shape needs no special-casing there; the legacy flat shape may still + carry a v0.2 `expert_hidden_dim`/`expert_n_blocks` override, honoured + only when that checkpoint's router was actually enabled. + + Routed models spend their FLOPs in the (smaller) expert trunks. Training + runs the full soft mixture (every expert on the whole batch), so its + activation memory scales with the expert count; inference does top-1 + dispatch (each row hits one expert), so the batch just partitions across + experts and one expert's dims already bound it. estimate_batch_size + scales memory linearly with hidden_dim * n_blocks, so the training + multiplier folds into n_blocks. """ - router_cfg = model_cfg.get("router") - if router_cfg and router_cfg.get("enabled"): - hidden_dim, n_blocks = gconfig.resolve_expert_dims( - router_cfg, model_cfg["hidden_dim"], model_cfg["n_blocks"] - ) - if training: - n_blocks *= _router_total_experts(router_cfg) - return hidden_dim, n_blocks - return model_cfg["hidden_dim"], model_cfg["n_blocks"] + if f"{stage}_model" in model_cfg: + stage_cfg = model_cfg[f"{stage}_model"] + hidden_dim, n_blocks = stage_cfg["hidden_dim"], stage_cfg["n_res_blocks"] + router_cfg = stage_cfg.get("router") + else: + hidden_dim, n_blocks = model_cfg["hidden_dim"], model_cfg["n_blocks"] + router_cfg = model_cfg.get("router") + if router_cfg and router_cfg.get("enabled"): + hidden_dim = model_cfg.get("expert_hidden_dim") or hidden_dim + n_blocks = model_cfg.get("expert_n_blocks") or n_blocks + if router_cfg and router_cfg.get("enabled") and training: + n_blocks = n_blocks * _router_total_experts(router_cfg) + return hidden_dim, n_blocks def _coerce_scalar(value: str) -> object: @@ -494,7 +520,6 @@ def train( cli_train = { k: v for k, v in { - "mode": mode.value if mode is not None else None, "epochs": epochs, "batch_size": batch_size_value, "lr": lr, @@ -507,9 +532,6 @@ def train( "validate_every": validate_every, "validate_steps": validate_steps, "max_val_batches": max_val_batches, - "n_critic": n_critic, - "gp_weight": gp_weight, - "critic_lr": critic_lr, "wandb": wandb, "wandb_project": wandb_project, "wandb_run_name": wandb_run_name, @@ -517,30 +539,79 @@ def train( }.items() if v is not None } - cli_model: dict[str, object] = { + # Architecture-shape flags apply to stage1_model only (decision: stage 2 + # is tuned via a hand-edited config.toml until stage-prefixed --stage2-* + # flags land — see docs/v0.3.0-design.md decision 7). + cli_stage1_model: dict[str, object] = { k: v for k, v in { "hidden_dim": hidden_dim, - "n_blocks": n_blocks, - "emb_dim": emb_dim, + "n_res_blocks": n_blocks, "dropout": dropout, - "conditioning": conditioning.value if conditioning is not None else None, - "noise_dim": noise_dim, }.items() if v is not None } cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) if cli_router: - cli_model["router"] = cli_router - cfg = gconfig.merge_cli_overrides( - gconfig.DEFAULT_CONFIG, config, cli_train, cli_model - ) - t, m = cfg["train"], cfg["model"] + cli_stage1_model["router"] = cli_router + + # --emb-dim/--conditioning set both conditioning axes (v0.2 had one + # shared value for particle+material). + cli_conditioning: dict[str, dict] = {} + if emb_dim is not None: + cli_conditioning["particle"] = {"emb_dim": emb_dim} + cli_conditioning["material"] = {"emb_dim": emb_dim} + if conditioning is not None: + cli_conditioning.setdefault("particle", {})["type"] = conditioning.value + cli_conditioning.setdefault("material", {})["type"] = conditioning.value + + overrides: dict[str, dict] = {} + if cli_train: + overrides["train"] = cli_train + if cli_stage1_model: + overrides["stage1_model"] = cli_stage1_model + if cli_conditioning: + overrides["conditioning"] = cli_conditioning + + # --mode/--n-critic/--gp-weight/--critic-lr apply to BOTH stages — v0.2 + # had one global mode/wgan config shared by both (see + # giant.config.migrate_config's train.mode / + # train.{n_critic,gp_weight,critic_lr} precedent), unlike the + # architecture-shape flags above. + if mode is not None: + overrides.setdefault("stage1_model", {})["generator"] = mode.value + overrides.setdefault("stage2_model", {})["generator"] = mode.value + if noise_dim is not None: + overrides.setdefault("stage1_model", {}).setdefault("wgan", {})["noise_dim"] = ( + noise_dim + ) + wgan_overrides = { + k: v + for k, v in { + "n_critic": n_critic, + "gp_weight": gp_weight, + "critic_lr": critic_lr, + }.items() + if v is not None + } + if wgan_overrides: + overrides.setdefault("stage1_model", {}).setdefault("wgan", {}).update( + wgan_overrides + ) + overrides.setdefault("stage2_model", {}).setdefault("wgan", {}).update( + wgan_overrides + ) + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides) + gconfig.validate_config(cfg) + t = cfg["train"] _device = torch.device(device) if device else gconfig.auto_device() if batch_size_auto: - est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m, training=True) + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims( + cfg, training=True, stage="stage1" + ) try: t["batch_size"] = gconfig.estimate_batch_size( est_hidden_dim, est_n_blocks, _device @@ -658,31 +729,45 @@ def new_run( cli_train = { k: v for k, v in { - "mode": mode.value if mode is not None else None, "epochs": epochs, "batch_size": batch_size, "lr": lr, }.items() if v is not None } - cli_model: dict[str, object] = { + cli_stage1_model: dict[str, object] = { k: v for k, v in { "hidden_dim": hidden_dim, - "n_blocks": n_blocks, - "emb_dim": emb_dim, + "n_res_blocks": n_blocks, "dropout": dropout, - "conditioning": conditioning.value if conditioning is not None else None, }.items() if v is not None } cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) if cli_router: - cli_model["router"] = cli_router + cli_stage1_model["router"] = cli_router + cli_conditioning: dict[str, dict] = {} + if emb_dim is not None: + cli_conditioning["particle"] = {"emb_dim": emb_dim} + cli_conditioning["material"] = {"emb_dim": emb_dim} + if conditioning is not None: + cli_conditioning.setdefault("particle", {})["type"] = conditioning.value + cli_conditioning.setdefault("material", {})["type"] = conditioning.value - cfg = gconfig.merge_cli_overrides( - gconfig.DEFAULT_CONFIG, config, cli_train, cli_model - ) + overrides: dict[str, dict] = {} + if cli_train: + overrides["train"] = cli_train + if cli_stage1_model: + overrides["stage1_model"] = cli_stage1_model + if cli_conditioning: + overrides["conditioning"] = cli_conditioning + if mode is not None: + overrides.setdefault("stage1_model", {})["generator"] = mode.value + overrides.setdefault("stage2_model", {})["generator"] = mode.value + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides) + gconfig.validate_config(cfg) run_dir = (out or gconfig.resolve_default_out_dir(cfg)).resolve() if not force: @@ -699,7 +784,7 @@ def new_run( if dry_run: typer.echo("dry-run: not writing anything. Resolved config:") - for section in ("train", "model"): + for section in ("train", "conditioning", "stage1_model", "stage2_model"): typer.echo(f"[{section}]") for k, v in cfg[section].items(): if k == "router": @@ -708,6 +793,11 @@ def new_run( return meta = { + # Tags the written config.toml as v0.3-shaped so a later + # `migrate_config` load (e.g. `giant train --config ...`) treats it + # as already-migrated instead of misreading it as v0.2 and dropping + # its stage1_model/stage2_model/conditioning content. + "config_version": gconfig.CONFIG_VERSION, "git_hash": gconfig.git_hash(), "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "created_by": "giant new-run", @@ -861,14 +951,22 @@ def predict( assert batch_size_value is not None bs = batch_size_value - conditioning = model_cfg.get("conditioning", "embedding") + conditioning = _conditioning_str(model_cfg) pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()} mat_map = {str(k): v for k, v in ckpt["mat_map"].items()} cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"]) - model, sec_decoder = build_models(model_cfg) + built = build_models(model_cfg) + model, sec_decoder = built["stage1"], built["stage2"] + if model is None or sec_decoder is None: + typer.echo( + "error: checkpoint has an inactive stage1 or stage2 — giant " + "predict needs both (see stage{1,2}_model.active)", + err=True, + ) + raise typer.Exit(1) _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) model.to(_device).eval() sec_decoder.to(_device).eval() @@ -1244,14 +1342,22 @@ def rollout( training_cfg = gconfig.load_checkpoint_config(checkpoint) model_cfg = ckpt["model_config"] - conditioning = model_cfg.get("conditioning", "embedding") + conditioning = _conditioning_str(model_cfg) pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()} mat_map = {str(k): v for k, v in ckpt["mat_map"].items()} cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"]) - model, sec_decoder = build_models(model_cfg) + built = build_models(model_cfg) + model, sec_decoder = built["stage1"], built["stage2"] + if model is None or sec_decoder is None: + typer.echo( + "error: checkpoint has an inactive stage1 or stage2 — giant " + "rollout needs both (see stage{1,2}_model.active)", + err=True, + ) + raise typer.Exit(1) _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) model.to(_device).eval() sec_decoder.to(_device).eval() diff --git a/giant/config.py b/giant/config.py index 77814af..1cb44f3 100644 --- a/giant/config.py +++ b/giant/config.py @@ -680,6 +680,18 @@ def validate_config(cfg: dict) -> None: single block's defaults. """ particle_type = _get_path(cfg, "conditioning.particle.type") + material_type = _get_path(cfg, "conditioning.material.type") + + if particle_type != material_type: + raise ValueError( + "conditioning.particle.type " + f"({particle_type!r}) != conditioning.material.type " + f"({material_type!r}) — the data pipeline (giant/data/transforms.py's " + "build_features/build_cond_features) doesn't support mixed " + "particle/material conditioning types yet, even though " + "ConditionEncoder itself (giant/model/network.py) is already able " + "to. Not implemented in v0.3.0 — use the same type for both axes." + ) pt_target = _get_path(cfg, "stage2_model.particle_type.target") if pt_target == "embedding" and particle_type != "embedding": diff --git a/giant/pipeline.py b/giant/pipeline.py index 21d6f6d..dfb2111 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -9,10 +9,8 @@ from torch.utils.data import DataLoader from giant import config from giant.constants import ( COND_DIM, - EMB_DIM, K_MAX, PARTICLE_PHYS_DIM, - SEC_SLOT_DIM, X_DIM, ) from giant.data import setup_cache @@ -56,26 +54,69 @@ class SetupStageResult: n_train_steps: int +def _seed_energy_router( + router_cfg: dict, + cond_norm: Normalizer, + energy_quantiles: np.ndarray, + energy_idx: int, + echo, +) -> None: + """Mutate `router_cfg["centers_init"]` in place from real data quantiles, + when this stage's router is an enabled EnergyRouter. Shared by both + stages' router configs — each seeded independently, since v0.3.0 stages + may have entirely different router configs (see docs/v0.3.0-design.md).""" + active = router_cfg.get("enabled") and router_cfg.get("type") == "energy" + if not active: + return + if energy_quantiles.size == 0: + echo( + " warning: no energy samples collected — EnergyRouter falls back to " + "default centers" + ) + return + assert cond_norm.mean is not None and cond_norm.std is not None + levels = np.linspace(0.0, 1.0, router_cfg["n_experts"]) + raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels) + centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[ + energy_idx + ] + router_cfg["centers_init"] = centers_init.astype(np.float32).tolist() + echo( + f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}" + ) + + def run_setup_stage( data: str | Path, val_fraction: float, seed: int, - conditioning: str, - router_cfg: dict, + cfg: dict, cache_setup: bool = True, rebuild_setup_cache: bool = False, echo=print, ) -> SetupStageResult: """Scan `data` for everything training needs before the epoch loop: the train/val event split, pdg/material vocab maps, an optional process map - (`router_cfg.type == "process"`), and the Stage-1/Stage-2 normalizers. + (needed if either stage's router is type="process"), and the Stage-1/ + Stage-2 normalizers. + + `cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/ + `stage2_model`), already passed through `giant.config.validate_config` — + in particular, `conditioning.particle.type == conditioning.material.type` + is assumed (the data pipeline doesn't support mixed conditioning types + yet, see `validate_config`), so a single shared conditioning string + drives `giant.data.transforms.build_features`. Reads from and writes to the `giant.data.setup_cache` sidecar when `cache_setup` is set (`rebuild_setup_cache` ignores — but still - refreshes — any existing sidecar content). `router_cfg` may be mutated - in place: an active `EnergyRouter` (`router_cfg["type"] == "energy"`) + refreshes — any existing sidecar content). Each stage's `router` config + is mutated in place: an active `EnergyRouter` (`router.type == "energy"`) gets its `centers_init` seeded from real data quantiles here. """ + conditioning = cfg["conditioning"]["particle"]["type"] + stage1_router = cfg["stage1_model"].get("router") or {} + stage2_router = cfg["stage2_model"].get("router") or {} + files = find_parquet_files(data) echo(f"found {len(files)} parquet file(s)") @@ -121,9 +162,23 @@ def run_setup_stage( if cache is not None: cache.vocab = (pdg_map, mat_map) + # A process map is needed if either stage's router reads the physics + # process label (type="process"). Only one map is built even if both + # stages want one — see the module-level note in giant/cli.py's + # _router_total_experts for why composed-router n_experts isn't a plain + # int; process routers are never composed in practice, so this doesn't + # need that generality. proc_map: dict[str, int] | None = None - if router_cfg.get("enabled") and router_cfg.get("type") == "process": - n_experts = router_cfg["n_experts"] + process_router_cfg = next( + ( + r + for r in (stage1_router, stage2_router) + if r.get("enabled") and r.get("type") == "process" + ), + None, + ) + if process_router_cfg is not None: + n_experts = process_router_cfg["n_experts"] cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None if cached_proc_map is not None: proc_map = cached_proc_map @@ -138,10 +193,11 @@ def run_setup_stage( if cache is not None: cache.proc_maps[n_experts] = proc_map - energy_router_active = ( - router_cfg.get("enabled") and router_cfg.get("type") == "energy" + energy_router_active = any( + r.get("enabled") and r.get("type") == "energy" + for r in (stage1_router, stage2_router) ) - energy_idx = router_cfg.get("energy_idx", 3) + energy_idx = 3 norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning) entry = cache.normalizers.get(norm_key) if cache is not None else None @@ -163,9 +219,9 @@ def run_setup_stage( # grid (setup_cache.energy_quantiles_from_sample) so centers can # instead be seeded from actual data quantiles below. Collected # whenever the setup cache is being populated, not only when *this* - # run's router is energy-typed, so a later run enabling - # --router-type energy against this same (val_fraction, seed, - # conditioning) key never needs to rescan just to seed centers. + # run's router is energy-typed, so a later run enabling an energy + # router against this same (val_fraction, seed, conditioning) key + # never needs to rescan just to seed centers. collect_energy_sample = energy_router_active or cache is not None energy_sampler = ( _ReservoirSampler(capacity=100_000) if collect_energy_sample else None @@ -206,22 +262,8 @@ def run_setup_stage( cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles ) - if energy_router_active and energy_quantiles.size > 0: - assert cond_norm.mean is not None and cond_norm.std is not None - levels = np.linspace(0.0, 1.0, router_cfg["n_experts"]) - raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels) - centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[ - energy_idx - ] - router_cfg["centers_init"] = centers_init.astype(np.float32).tolist() - echo( - f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}" - ) - elif energy_router_active: - echo( - " warning: no energy samples collected — EnergyRouter falls back to " - "default centers" - ) + _seed_energy_router(stage1_router, cond_norm, energy_quantiles, energy_idx, echo) + _seed_energy_router(stage2_router, cond_norm, energy_quantiles, energy_idx, echo) if cache is not None: setup_cache.save(data, files, cache, echo=echo) @@ -252,7 +294,7 @@ def run_train_job( rebuild_setup_cache: bool = False, echo=print, ) -> None: - t, m = cfg["train"], cfg["model"] + t = cfg["train"] config.seed_everything(t["seed"]) out_dir = Path(out_dir) @@ -271,20 +313,13 @@ def run_train_job( "section)" ) - router_cfg = m["router"] - if t["mode"] == "wgan" and router_cfg.get("enabled"): - raise ValueError( - "--mode wgan does not support --router (no routed WGAN generator/" - "critic exists) — disable one or the other" - ) - - conditioning = m["conditioning"] + config.validate_config(cfg) + conditioning = cfg["conditioning"]["particle"]["type"] setup = run_setup_stage( data, val_fraction=t["val_fraction"], seed=t["seed"], - conditioning=conditioning, - router_cfg=router_cfg, + cfg=cfg, cache_setup=cache_setup, rebuild_setup_cache=rebuild_setup_cache, echo=echo, @@ -347,60 +382,19 @@ def run_train_job( pin_memory=pin, ) - emb_dim = m.get("emb_dim", EMB_DIM) - expert_hidden_dim, expert_n_blocks = config.resolve_expert_dims( - router_cfg, m["hidden_dim"], m["n_blocks"] - ) - if router_cfg.get("enabled") and (expert_hidden_dim, expert_n_blocks) != ( - m["hidden_dim"], - m["n_blocks"], - ): - # Only reachable via an explicit router.expert_hidden_dim/n_blocks - # override (the 0/"unset" sentinel always resolves to m["hidden_dim"]/ - # ["n_blocks"] — see resolve_expert_dims), so this is never a false - # positive from inheritance, only a deliberate narrow/wide-experts - # config the checkpoint dir name (_h{hidden_dim}_b{n_blocks}) won't - # reflect. - echo( - f" warning: experts are {expert_hidden_dim}x{expert_n_blocks}, " - f"different from model.hidden_dim/n_blocks ({m['hidden_dim']}x" - f"{m['n_blocks']}) — the checkpoint dir name reflects the latter, " - "not the experts actually being trained" - ) - model_config = { "pdg_vocab": len(pdg_map), "mat_vocab": len(mat_map), - "hidden_dim": m["hidden_dim"], - "n_blocks": m["n_blocks"], - "emb_dim": emb_dim, - "dropout": m["dropout"], - "k_max": K_MAX, - "sec_slot_dim": SEC_SLOT_DIM, - "conditioning": conditioning, - "router": dict(router_cfg), - "expert_hidden_dim": expert_hidden_dim, - "expert_n_blocks": expert_n_blocks, - # Read by `predict`/`rollout` (which never receive their own --mode - # flag) to auto-detect which sampler a checkpoint needs. - "mode": t["mode"], - "noise_dim": m.get("noise_dim", 64), + "conditioning": cfg["conditioning"], + "stage1_model": cfg["stage1_model"], + "stage2_model": cfg["stage2_model"], } - stage1_model, sec_decoder = build_models(model_config) - echo( - f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters | " - f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters" - ) - - critic = None - sec_critic = None - if t["mode"] == "wgan": - critic, sec_critic = build_critics(model_config) - echo( - f"critic: {sum(p.numel() for p in critic.parameters()):,} parameters | " - f"sec_critic: {sum(p.numel() for p in sec_critic.parameters()):,} parameters" - ) + models = build_models(model_config) + critics = build_critics(model_config) + for name, model in models.items(): + if model is not None: + echo(f"{name}: {sum(p.numel() for p in model.parameters()):,} parameters") out_dir.mkdir(parents=True, exist_ok=True) meta = config.build_run_meta( @@ -415,25 +409,13 @@ def run_train_job( config.save_config(cfg, out_dir, meta) run_training( - stage1_model=stage1_model, - sec_decoder=sec_decoder, + cfg=cfg, + models=models, + critics=critics, train_loader=train_loader, val_loader=val_loader, - mode=t["mode"], - epochs=t["epochs"], - lr=t["lr"], - weight_decay=t["weight_decay"], - ema_decay=t["ema_decay"], - warmup_epochs=t["warmup_epochs"], device=device, out_dir=out_dir, - lambda_nsec=t.get("lambda_nsec", 0.1), - lambda_s2=t.get("lambda_s2", 1.0), - lambda_balance=router_cfg.get("lambda_balance", 0.0), - lambda_proc=router_cfg.get("lambda_proc", 0.0), - lambda_entropy=router_cfg.get("lambda_entropy", 0.0), - gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0), - gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1), normalizer_dict={ "cond": cond_norm.to_dict(), "target": tgt_norm.to_dict(), @@ -444,16 +426,8 @@ def run_train_job( proc_map=proc_map, model_config=model_config, resume_path=resume, - validate_every=t["validate_every"], - validate_steps=t["validate_steps"], - max_val_batches=t["max_val_batches"], total_train_batches=total_train_batches, - critic=critic, - sec_critic=sec_critic, - n_critic=t.get("n_critic", 5), - gp_weight=t.get("gp_weight", 10.0), - critic_lr=t.get("critic_lr") or None, - use_wandb=t.get("wandb", False), + use_wandb=t.get("wandb", True), wandb_project=t.get("wandb_project", "giant"), wandb_run_name=t.get("wandb_run_name", ""), wandb_log_every=t.get("wandb_log_every", 50), diff --git a/giant/train.py b/giant/train.py index a6324bc..2e1688c 100644 --- a/giant/train.py +++ b/giant/train.py @@ -4,6 +4,7 @@ import math import os import signal import time +import warnings from pathlib import Path from types import FrameType from typing import Callable @@ -15,7 +16,8 @@ import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm -from giant.constants import K_MAX, SEC_SLOT_DIM +from giant.constants import K_MAX +from giant.model.network import Router from giant.model.schedule import ( CosineSchedule, flow_matching_loss, @@ -24,48 +26,6 @@ from giant.model.schedule import ( from giant.model.wgan import gradient_penalty, generator_loss from giant.validate import validate_marginals -_METRICS_FIELDS = [ - "epoch", - "train_loss", - "train_loss_s1", - "train_loss_nsec", - "train_loss_s2", - "train_loss_balance", - "train_loss_proc", - "train_loss_entropy", - "train_nsec_acc", - "d_loss", - "g_loss", - "wasserstein_estimate", - "gp_loss", - "val_loss", - "val_loss_s1", - "val_loss_nsec", - "val_loss_s2", - "val_loss_balance", - "val_loss_proc", - "val_loss_entropy", - "val_nsec_acc", - "val_marginal_kl", - "router_s1_entropy", - "router_s1_util_min", - "router_s1_util_max", - "router_s1_util_std", - "router_s2_entropy", - "router_s2_util_min", - "router_s2_util_max", - "router_s2_util_std", - "lr", - "critic_lr", - "grad_norm", - "grad_norm_d", - "grad_norm_g", - "gpu_mem_mb", - "samples_per_sec", - "is_best", - "epoch_time_s", -] - _CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM) @@ -126,342 +86,718 @@ def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) - return tau_start + (tau_end - tau_start) * progress -def _wandb_run_config( - *, - mode: str, - epochs: int, - lr: float, - warmup_epochs: int, - weight_decay: float, - ema_decay: float, - lambda_nsec: float, - lambda_s2: float, - lambda_balance: float, - lambda_proc: float, - lambda_entropy: float, - gumbel_tau_start: float, - gumbel_tau_end: float, - n_critic: int, - gp_weight: float, - model_config: dict | None, - stage1_params: int, - sec_decoder_params: int, - critic_params: int, - sec_critic_params: int, - total_params: int, -) -> dict: - """Build the dict logged as a wandb run's `config`. +def _stage_router(model: torch.nn.Module) -> Router | None: + """A stage model's Router, if its trunk is routed — else None. - Router-only knobs (`lambda_balance`/`lambda_proc`/`lambda_entropy`/ - `gumbel_tau_start`/`gumbel_tau_end`) and WGAN-only knobs (`n_critic`/ - `gp_weight`) are omitted unless actually active, so a run's wandb config - doesn't imply hyperparameters from an inactive code path (a disabled - router's fine-tuning knobs, or GAN critic settings for a flow/DDPM run). - The full `model_config` (including its `router` sub-dict, whatever the - router type/state) is always included, so no information is lost — this - only trims the flattened top-level convenience duplicates. + Post-step-2 refactor the router lives at `model.trunk.router` + (`giant.model.network.RoutedTrunk`), not `model.router` directly. """ - router_enabled = bool((model_config or {}).get("router", {}).get("enabled", False)) - cfg = { - "mode": mode, - "epochs": epochs, - "lr": lr, - "warmup_epochs": warmup_epochs, - "weight_decay": weight_decay, - "ema_decay": ema_decay, - "lambda_nsec": lambda_nsec, - "lambda_s2": lambda_s2, - "model": model_config or {}, - "stage1_params": stage1_params, - "sec_decoder_params": sec_decoder_params, - "critic_params": critic_params, - "sec_critic_params": sec_critic_params, - "total_params": total_params, - } - if router_enabled: - cfg.update( - { - "lambda_balance": lambda_balance, - "lambda_proc": lambda_proc, - "lambda_entropy": lambda_entropy, - "gumbel_tau_start": gumbel_tau_start, - "gumbel_tau_end": gumbel_tau_end, - } + trunk = getattr(model, "trunk", None) + return getattr(trunk, "router", None) + + +def _batch_to_device(batch: tuple, device: torch.device) -> tuple: + return tuple(t.to(device) for t in batch) + + +class StageTrainer: + """One active stage's optimizer(s), EMA, and per-batch step. + + Reads only the shared batch tuple `(cond_cont, cond_cat, x1_s1, n_sec, + sec_cont, proc_idx)` — stage 2 always conditions on the ground-truth + `x1_s1` (`stage2_model.stage1_context = "truth"`, stage-level teacher + forcing; `"sampled"` is not implemented — see docs/v0.3.0-design.md §3.3), + so stage trainers never need each other's output at train time. This + means "stage-2-only training is a cheap ablation, not new plumbing" + (design doc §7) falls out for free: a trainer only exists for active + stages, and inactive stages are simply never constructed. + + Grad-norm clipping is per-stage here — v0.2's single shared optimizer + clipped both stages' gradients jointly; splitting per stage is a small, + disclosed behavior change. It doesn't affect Adam's per-parameter update + math itself (no cross-parameter coupling), only the clip threshold's + scope. + """ + + name: str + is_stage2: bool + model: torch.nn.Module + ema_model: torch.nn.Module | None + router: Router | None + optimizer: optim.Optimizer + + def step(self, batch: tuple, device: torch.device, global_step: int) -> dict: + raise NotImplementedError + + def val_loss(self, batch: tuple, device: torch.device) -> dict: + raise NotImplementedError + + def train_mode(self) -> None: + raise NotImplementedError + + def eval_mode(self) -> None: + raise NotImplementedError + + def sampling_model(self) -> torch.nn.Module: + return self.ema_model if self.ema_model is not None else self.model + + def state_dict(self) -> dict: + raise NotImplementedError + + def load_state_dict(self, sd: dict) -> None: + raise NotImplementedError + + def resume_lr(self, lr: float, critic_lr: float = 0.0) -> None: + raise NotImplementedError + + +class FlowDDPMStageTrainer(StageTrainer): + """flow or ddpm generator for a single stage.""" + + def __init__( + self, + name: str, + model: torch.nn.Module, + is_stage2: bool, + generator: str, + lambda_weight: float, + n_sec_lambda: float, + lambda_balance: float, + lambda_proc: float, + lambda_entropy: float, + gumbel_tau_start: float, + gumbel_tau_end: float, + lr: float, + weight_decay: float, + ema_decay: float, + warmup_epochs: int, + epochs: int, + steps_per_epoch: int, + ddpm_n_steps: int, + device: torch.device, + ) -> None: + if is_stage2 and generator not in ("flow",): + raise NotImplementedError( + f"stage2_model.generator={generator!r} is not implemented for " + "stage 2 (only 'flow' and 'wgan' — 'ddpm' has no stage-2 " + "secondary-decoder loss yet, see docs/v0.3.0-design.md)" + ) + self.name = name + self.is_stage2 = is_stage2 + self.generator = generator + self.device = device + self.model = model.to(device) + self.lambda_weight = lambda_weight + self.n_sec_lambda = n_sec_lambda + self.lambda_balance = lambda_balance + self.lambda_proc = lambda_proc + self.lambda_entropy = lambda_entropy + self.gumbel_tau_start = gumbel_tau_start + self.gumbel_tau_end = gumbel_tau_end + self.ema_decay = ema_decay + self.router = _stage_router(self.model) + + self.params = list(self.model.parameters()) + self.optimizer = optim.AdamW(self.params, lr=lr, weight_decay=weight_decay) + warmup_steps = warmup_epochs * steps_per_epoch + total_steps = max(epochs * steps_per_epoch, 1) + + def _lr_lambda(step: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + return (step + 1) / warmup_steps + t = step - warmup_steps + T = max(total_steps - warmup_steps, 1) + return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T)) + + self._lr_lambda = _lr_lambda + self.lr_sched = optim.lr_scheduler.LambdaLR(self.optimizer, _lr_lambda) + self.total_steps = total_steps + self.ddpm_schedule = ( + CosineSchedule(T=ddpm_n_steps).to(device) if generator == "ddpm" else None ) - if mode == "wgan": - cfg.update({"n_critic": n_critic, "gp_weight": gp_weight}) - return cfg + + self.ema_model: torch.nn.Module | None = None + if ema_decay > 0: + self.ema_model = copy.deepcopy(self.model).eval() + for p in self.ema_model.parameters(): + p.requires_grad_(False) + + def _predict_n_sec(self, cond_cont, cond_cat, stage1_ctx): + if self.model.n_sec_head is None: + return None + if self.is_stage2: + return self.model.predict_n_sec(cond_cont, cond_cat, stage1_ctx) + return self.model.predict_n_sec(cond_cont, cond_cat) + + def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx): + if not self.is_stage2: + if self.generator == "flow": + return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat) + assert self.ddpm_schedule is not None + return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat) + return flow_matching_loss_secondary( + self.model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask + ) + + def _compute(self, batch: tuple, device: torch.device) -> dict: + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = _batch_to_device( + batch, device + ) + x1_s2 = sec_cont.flatten(1) + sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) + stage1_ctx = x1_s1.detach() + + l_gen = self._generator_loss( + cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx + ) + n_sec_logits = self._predict_n_sec(cond_cont, cond_cat, stage1_ctx) + l_nsec = torch.zeros((), device=device) + nsec_acc = torch.zeros((), device=device) + if n_sec_logits is not None: + l_nsec = F.cross_entropy(n_sec_logits, n_sec) + nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean() + + l_balance = l_proc = l_entropy = torch.zeros((), device=device) + if self.router is not None: + l_balance = self.router.balance_loss(cond_cont, cond_cat) + l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx) + l_entropy = self.router.entropy_loss(cond_cont, cond_cat) + + total = self.lambda_weight * l_gen + self.n_sec_lambda * l_nsec + if self.lambda_balance > 0: + total = total + self.lambda_balance * l_balance + if self.lambda_proc > 0: + total = total + self.lambda_proc * l_proc + if self.lambda_entropy > 0: + total = total + self.lambda_entropy * l_entropy + + return { + "total": total, + "loss_gen": l_gen, + "loss_nsec": l_nsec, + "loss_balance": l_balance, + "loss_proc": l_proc, + "loss_entropy": l_entropy, + "nsec_acc": nsec_acc, + "batch_size": cond_cont.size(0), + } + + def step(self, batch: tuple, device: torch.device, global_step: int) -> dict: + if self.router is not None: + self.router.gumbel_tau = _gumbel_tau( + global_step, + self.total_steps, + self.gumbel_tau_start, + self.gumbel_tau_end, + ) + out = self._compute(batch, device) + self.optimizer.zero_grad() + out["total"].backward() + grad_norm = torch.nn.utils.clip_grad_norm_(self.params, 1.0) + self.optimizer.step() + self.lr_sched.step() + if self.ema_model is not None: + _update_ema(self.ema_model, self.model, self.ema_decay) + return { + "loss": out["total"].item(), + "loss_gen": out["loss_gen"].item(), + "loss_nsec": out["loss_nsec"].item(), + "loss_balance": out["loss_balance"].item(), + "loss_proc": out["loss_proc"].item(), + "loss_entropy": out["loss_entropy"].item(), + "nsec_acc": out["nsec_acc"].item(), + "grad_norm": grad_norm.item(), + "lr": self.optimizer.param_groups[0]["lr"], + "batch_size": out["batch_size"], + } + + @torch.no_grad() + def val_loss(self, batch: tuple, device: torch.device) -> dict: + out = self._compute(batch, device) + return { + "loss": out["total"].item(), + "loss_gen": out["loss_gen"].item(), + "loss_nsec": out["loss_nsec"].item(), + "loss_balance": out["loss_balance"].item(), + "loss_proc": out["loss_proc"].item(), + "loss_entropy": out["loss_entropy"].item(), + "nsec_acc": out["nsec_acc"].item(), + "batch_size": out["batch_size"], + } + + def train_mode(self) -> None: + self.model.train() + + def eval_mode(self) -> None: + self.model.eval() + + def state_dict(self) -> dict: + sd = { + "model": self.model.state_dict(), + "optimizer": self.optimizer.state_dict(), + "lr_sched": self.lr_sched.state_dict(), + } + if self.ema_model is not None: + sd["model_ema"] = self.ema_model.state_dict() + return sd + + def load_state_dict(self, sd: dict) -> None: + self.model.load_state_dict(sd["model"]) + self.optimizer.load_state_dict(sd["optimizer"]) + self.lr_sched.load_state_dict(sd["lr_sched"]) + if self.ema_model is not None: + self.ema_model.load_state_dict(sd.get("model_ema", sd["model"])) + + def resume_lr(self, lr: float, critic_lr: float = 0.0) -> None: + """Restore --lr's authority after load_state_dict restored the + checkpoint's own base LR (mirrors the resume fixup train() used to + do inline). `critic_lr` is unused here (no critic on this trainer), + kept only to match `StageTrainer.resume_lr`'s signature.""" + self.lr_sched.base_lrs = [lr for _ in self.lr_sched.base_lrs] + resumed_lr = lr * self._lr_lambda(self.lr_sched.last_epoch) + for group in self.optimizer.param_groups: + group["lr"] = resumed_lr -def _compute_losses( - stage1_model: torch.nn.Module, - sec_decoder: torch.nn.Module, - batch: tuple, - mode: str, - ddpm_schedule, - device: torch.device, - lambda_nsec: float, - lambda_s2: float, - lambda_balance: float = 0.0, - lambda_proc: float = 0.0, - lambda_entropy: float = 0.0, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, L_entropy, nsec_acc) for one batch.""" - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch - cond_cont = cond_cont.to(device) - cond_cat = cond_cat.to(device) - x1_s1 = x1_s1.to(device) - n_sec = n_sec.to(device) - sec_cont = sec_cont.to(device) - proc_idx = proc_idx.to(device) +class WGANStageTrainer(StageTrainer): + """WGAN-GP generator+critic for a single stage (see giant/model/wgan.py). - # Stage-1 flow loss - if mode == "flow": - l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat) - else: - assert ddpm_schedule is not None - l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat) - - # n_sec classification loss - n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat) - l_nsec = F.cross_entropy(n_sec_logits, n_sec) - nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean() - - # Stage-2 secondary flow loss - # Use a noiseless Stage-1 target as context (detach to avoid back-prop - # coupling between the two flow paths). sec_cont's log_mass/charge - # columns are already a fixed physics-derived regression target (see - # giant.data.transforms.encode_secondaries) rather than a learned/moving - # one, so — unlike the embedding-table target this replaced — no - # detaching is needed to keep the target from chasing the decoder. - x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM) - - sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) - l_s2 = flow_matching_loss_secondary( - sec_decoder, - x1_s2, - cond_cont, - cond_cat, - x1_s1.detach(), - sec_mask, - ) - - # Optional MoE load-balance auxiliary loss: only present when both stages - # are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`, - # the monolith models don't), computed on cond_cont alone (cheap — no - # trunk compute) so it's reported even when lambda_balance == 0. - if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"): - l_balance = stage1_model.router.balance_loss( - cond_cont, cond_cat - ) + sec_decoder.router.balance_loss(cond_cont, cond_cat) - # Supervised router auxiliary loss (e.g. ProcessRouter's process - # classifier); a scalar 0 for routers with no such loss (EnergyRouter). - l_proc = stage1_model.router.classify_loss( - cond_cont, cond_cat, proc_idx - ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx) - # Optional entropy-regularization aux loss (see Router.entropy_loss): - # penalizes uniform/collapsed gating, a secondary guard against - # gate-sharpness collapse that lambda_balance alone can't see. - l_entropy = stage1_model.router.entropy_loss( - cond_cont, cond_cat - ) + sec_decoder.router.entropy_loss(cond_cont, cond_cat) - else: - l_balance = torch.zeros((), device=device) - l_proc = torch.zeros((), device=device) - l_entropy = torch.zeros((), device=device) - - total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2 - if lambda_balance > 0: - total = total + lambda_balance * l_balance - if lambda_proc > 0: - total = total + lambda_proc * l_proc - if lambda_entropy > 0: - total = total + lambda_entropy * l_entropy - return total, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc - - -def _wgan_train_step( - generator: torch.nn.Module, - sec_generator: torch.nn.Module, - critic: torch.nn.Module, - sec_critic: torch.nn.Module, - batch: tuple, - device: torch.device, - optimizer_g: optim.Optimizer, - optimizer_d: optim.Optimizer, - g_params: list, - d_params: list, - step_count: int, - n_critic: int, - gp_weight: float, - lambda_nsec: float, - lambda_s2: float, -) -> dict: - """One WGAN-GP training step, both stages (see giant/model/wgan.py for the losses). - - Both critics update every batch. Every `n_critic`-th batch additionally - updates both generators. The (non-adversarial) n_sec classifier updates - every batch regardless — folded into whichever `optimizer_g` step happens - this batch (full adversarial g_loss on generator batches, n_sec-only in - between) rather than throttled to the generator's cadence, since n_sec - accuracy is a headline flow-vs-wgan comparison metric and shares the - generator's ConditionEncoder. - - Stage 2's real/fake target is a flattened (B, SEC_DIM) vector with - `K_MAX - n_sec` padded slots per row; both critic's input and its - gradient-penalty gradient are masked to the valid slots (see - `giant.model.wgan.gradient_penalty`) so the critic can't key on padding - instead of genuine content. Stage 2 is conditioned on the *real* - ground-truth Stage-1 target (`x1_s1`, detached) rather than the - generator's own fake Stage-1 output — same precedent as the flow-matching - path's `flow_matching_loss_secondary` call, avoiding compounding errors - during training. + Ports `_wgan_train_step` to operate on one stage instead of two fused + together — the critic updates every batch; every `n_critic`-th batch + additionally updates the generator (`did_g_step`). The (non-adversarial) + n_sec classifier, when this stage's model owns it, updates every batch + regardless — folded into whichever generator optimizer step happens this + batch, same precedent as v0.2. """ - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, _proc_idx = batch - cond_cont = cond_cont.to(device) - cond_cat = cond_cat.to(device) - x1_s1 = x1_s1.to(device) - n_sec = n_sec.to(device) - sec_cont = sec_cont.to(device) - B = x1_s1.size(0) - x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM) - sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) - mask_flat = ( - sec_mask.unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1).float() - ) - stage1_ctx = x1_s1.detach() + def __init__( + self, + name: str, + model: torch.nn.Module, + critic: torch.nn.Module, + is_stage2: bool, + lambda_weight: float, + n_sec_lambda: float, + n_critic: int, + gp_weight: float, + lr: float, + critic_lr: float, + ema_decay: float, + warmup_epochs: int, + epochs: int, + steps_per_epoch: int, + device: torch.device, + ) -> None: + self.name = name + self.is_stage2 = is_stage2 + self.device = device + self.model = model.to(device) + self.critic = critic.to(device) + self.lambda_weight = lambda_weight + self.n_sec_lambda = n_sec_lambda + self.n_critic = max(n_critic, 1) + self.gp_weight = gp_weight + self.ema_decay = ema_decay + self.router = _stage_router(self.model) - def critic_fn1(x: torch.Tensor) -> torch.Tensor: - return critic(x, cond_cont, cond_cat) + self.g_params = list(self.model.parameters()) + self.d_params = list(self.critic.parameters()) + # WGAN-GP recipe (Gulrajani et al. 2017): Adam, beta1=0, no weight decay. + self.optimizer = optim.Adam(self.g_params, lr=lr, betas=(0.0, 0.9)) + self.optimizer_d = optim.Adam( + self.d_params, lr=critic_lr if critic_lr > 0 else lr, betas=(0.0, 0.9) + ) - def critic_fn2(x: torch.Tensor) -> torch.Tensor: - return sec_critic(x, cond_cont, cond_cat, stage1_ctx) + # Generator steps fire every n_critic-th batch, so warmup/decay must + # be counted in those units, matching v0.2. + gen_steps_per_epoch = max(steps_per_epoch // self.n_critic, 1) + warmup_steps = warmup_epochs * gen_steps_per_epoch + total_steps = max(epochs * gen_steps_per_epoch, 1) - z1 = torch.randn(B, generator.noise_dim, device=device) - fake1 = generator(z1, cond_cont, cond_cat) - z2 = torch.randn(B, sec_generator.noise_dim, device=device) - fake2 = sec_generator(z2, cond_cont, cond_cat, stage1_ctx) - fake2_masked = fake2 * mask_flat - real2_masked = x1_s2 * mask_flat + def _lr_lambda(step: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + return (step + 1) / warmup_steps + t = step - warmup_steps + T = max(total_steps - warmup_steps, 1) + return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T)) - # --- Critic step (every batch) --- - fake1_detached = fake1.detach() - real1_score = critic_fn1(x1_s1) - fake1_score = critic_fn1(fake1_detached) - gp1 = gradient_penalty(critic_fn1, x1_s1, fake1_detached) - d1 = fake1_score.mean() - real1_score.mean() + gp_weight * gp1 - wasserstein_estimate = (real1_score.mean() - fake1_score.mean()).detach() + self._lr_lambda = _lr_lambda + self.lr_sched = optim.lr_scheduler.LambdaLR(self.optimizer, _lr_lambda) - fake2_detached_masked = fake2_masked.detach() - real2_score = critic_fn2(real2_masked) - fake2_score = critic_fn2(fake2_detached_masked) - gp2 = gradient_penalty( - critic_fn2, real2_masked, fake2_detached_masked, mask=mask_flat - ) - d2 = fake2_score.mean() - real2_score.mean() + gp_weight * gp2 + self.ema_model: torch.nn.Module | None = None + if ema_decay > 0: + self.ema_model = copy.deepcopy(self.model).eval() + for p in self.ema_model.parameters(): + p.requires_grad_(False) - d_loss = d1 + lambda_s2 * d2 - optimizer_d.zero_grad() - d_loss.backward() - grad_norm_d = torch.nn.utils.clip_grad_norm_(d_params, 1.0) - optimizer_d.step() + def step(self, batch: tuple, device: torch.device, global_step: int) -> dict: + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, _proc_idx = _batch_to_device( + batch, device + ) + B = cond_cont.size(0) + stage1_ctx = x1_s1.detach() - # --- Generator (+ n_sec) step --- - did_g_step = step_count % n_critic == 0 - n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat) - l_nsec = F.cross_entropy(n_sec_logits, n_sec) - nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean() - optimizer_g.zero_grad() - if did_g_step: - g1 = generator_loss(critic_fn1, fake1) - g2 = generator_loss(critic_fn2, fake2_masked) - g_loss = g1 + lambda_nsec * l_nsec + lambda_s2 * g2 - else: - g1 = torch.zeros((), device=device) - g2 = torch.zeros((), device=device) - g_loss = lambda_nsec * l_nsec - g_loss.backward() - grad_norm_g = torch.nn.utils.clip_grad_norm_(g_params, 1.0) - optimizer_g.step() + if not self.is_stage2: + real = x1_s1 + def critic_fn(x): + return self.critic(x, cond_cont, cond_cat) + + z = torch.randn(B, self.model.noise_dim, device=device) + fake = self.model(z, cond_cont, cond_cat) + mask = None + else: + from giant.constants import SEC_SLOT_DIM + + sec_mask = torch.arange(K_MAX, device=device).unsqueeze( + 0 + ) < n_sec.unsqueeze(1) + mask = ( + sec_mask.unsqueeze(-1) + .expand(-1, -1, SEC_SLOT_DIM) + .reshape(B, -1) + .float() + ) + real = sec_cont.flatten(1) * mask + + def critic_fn(x): + return self.critic(x, cond_cont, cond_cat, stage1_ctx) + + z = torch.randn(B, self.model.noise_dim, device=device) + fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx) + fake = fake_raw * mask + + # --- critic step (every batch) --- + fake_detached = fake.detach() + real_score = critic_fn(real) + fake_score = critic_fn(fake_detached) + gp = gradient_penalty(critic_fn, real, fake_detached, mask=mask) + d_loss = fake_score.mean() - real_score.mean() + self.gp_weight * gp + wasserstein_estimate = (real_score.mean() - fake_score.mean()).detach() + + self.optimizer_d.zero_grad() + d_loss.backward() + grad_norm_d = torch.nn.utils.clip_grad_norm_(self.d_params, 1.0) + self.optimizer_d.step() + + # --- generator (+ n_sec) step --- + did_g_step = global_step % self.n_critic == 0 + n_sec_logits = None + if self.model.n_sec_head is not None: + n_sec_logits = ( + self.model.predict_n_sec(cond_cont, cond_cat) + if not self.is_stage2 + else self.model.predict_n_sec(cond_cont, cond_cat, stage1_ctx) + ) + l_nsec = torch.zeros((), device=device) + nsec_acc = torch.zeros((), device=device) + if n_sec_logits is not None: + l_nsec = F.cross_entropy(n_sec_logits, n_sec) + nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean() + + # On a non-generator-step batch with no n_sec_head on this stage + # (n_sec now defaults to stage 2, decision 1), there's nothing for + # the generator optimizer to do this batch — g_loss would otherwise + # be a graph-less zero tensor, which .backward() rejects outright. + skip_g_step = not did_g_step and n_sec_logits is None + if did_g_step: + g_loss_adv = generator_loss(critic_fn, fake) + g_loss = self.lambda_weight * g_loss_adv + self.n_sec_lambda * l_nsec + else: + g_loss_adv = torch.zeros((), device=device) + g_loss = self.n_sec_lambda * l_nsec + if skip_g_step: + grad_norm_g = torch.zeros(()) + else: + self.optimizer.zero_grad() + g_loss.backward() + grad_norm_g = torch.nn.utils.clip_grad_norm_(self.g_params, 1.0) + self.optimizer.step() + + if did_g_step: + self.lr_sched.step() + if self.ema_model is not None: + _update_ema(self.ema_model, self.model, self.ema_decay) + + return { + "d_loss": d_loss.item(), + "g_loss": g_loss_adv.item(), + "wasserstein_estimate": wasserstein_estimate.item(), + "gp_loss": gp.item(), + "loss_nsec": l_nsec.item(), + "nsec_acc": nsec_acc.item(), + "did_g_step": did_g_step, + "grad_norm": grad_norm_d.item() + grad_norm_g.item(), + "grad_norm_d": grad_norm_d.item(), + "grad_norm_g": grad_norm_g.item(), + "lr": self.optimizer.param_groups[0]["lr"], + "critic_lr": self.optimizer_d.param_groups[0]["lr"], + "batch_size": B, + } + + def val_loss(self, batch: tuple, device: torch.device) -> dict: + # No monotone per-batch WGAN loss fit for averaging; best-checkpoint + # selection instead uses validate_marginals (or its fallback) — see + # train()'s end-of-epoch block. + raise NotImplementedError + + def train_mode(self) -> None: + self.model.train() + self.critic.train() + + def eval_mode(self) -> None: + self.model.eval() + self.critic.eval() + + def state_dict(self) -> dict: + sd = { + "model": self.model.state_dict(), + "critic": self.critic.state_dict(), + "optimizer": self.optimizer.state_dict(), + "optimizer_d": self.optimizer_d.state_dict(), + "lr_sched": self.lr_sched.state_dict(), + } + if self.ema_model is not None: + sd["model_ema"] = self.ema_model.state_dict() + return sd + + def load_state_dict(self, sd: dict) -> None: + self.model.load_state_dict(sd["model"]) + self.critic.load_state_dict(sd["critic"]) + self.optimizer.load_state_dict(sd["optimizer"]) + self.optimizer_d.load_state_dict(sd["optimizer_d"]) + self.lr_sched.load_state_dict(sd["lr_sched"]) + if self.ema_model is not None: + self.ema_model.load_state_dict(sd.get("model_ema", sd["model"])) + + def resume_lr(self, lr: float, critic_lr: float = 0.0) -> None: + self.lr_sched.base_lrs = [lr for _ in self.lr_sched.base_lrs] + resumed_lr = lr * self._lr_lambda(self.lr_sched.last_epoch) + for group in self.optimizer.param_groups: + group["lr"] = resumed_lr + resumed_critic_lr = critic_lr if critic_lr > 0 else lr + for group in self.optimizer_d.param_groups: + group["lr"] = resumed_critic_lr + + +_WARNED_MARGINAL_VALIDATION_BROKEN = False + + +def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs): + """`validate_marginals` delegates to `giant.sample`'s samplers, which + still assume Stage 1 always owns `n_sec_head` (v0.2 behaviour) — under + decision 1 (docs/v0.3.0-design.md §2) a fresh Stage2OneShot owns it by + default instead, so this currently raises for any non-legacy checkpoint. + `giant/sample.py` needs a per-stage generator/n_sec update (design doc + §10, deferred to step 6); until then this degrades gracefully with a + one-time warning instead of crashing the whole training run. + """ + global _WARNED_MARGINAL_VALIDATION_BROKEN + model = trainer.sampling_model() + try: + return validate_marginals(model, val_loader, device=device, **kwargs) + except Exception as exc: # noqa: BLE001 — see docstring: any failure here + # is expected until giant/sample.py's per-stage n_sec update lands. + if not _WARNED_MARGINAL_VALIDATION_BROKEN: + warnings.warn( + "marginal validation unavailable this run " + f"({type(exc).__name__}: {exc}) — giant/sample.py doesn't yet " + "support a Stage2-owned n_sec head (docs/v0.3.0-design.md " + "step 6); val_marginal_kl stays NaN, and wgan best-checkpoint " + "selection falls back to the Wasserstein-distance magnitude.", + stacklevel=2, + ) + _WARNED_MARGINAL_VALIDATION_BROKEN = True + return None + + +def _build_stage_trainers( + cfg: dict, + models: dict[str, torch.nn.Module | None], + critics: dict[str, torch.nn.Module | None], + device: torch.device, + total_train_batches: int, +) -> dict[str, StageTrainer]: + t = cfg["train"] + trainers: dict[str, StageTrainer] = {} + for name, is_stage2 in (("stage1", False), ("stage2", True)): + model = models.get(name) + if model is None: + continue + stage_cfg = cfg[f"{name}_model"] + generator = stage_cfg["generator"] + lambda_weight = stage_cfg.get("lambda", 1.0) + n_sec_cfg = cfg["stage2_model"].get("n_sec", {}) + n_sec_lambda = n_sec_cfg.get("lambda", 0.1) + router_cfg = stage_cfg.get("router") or {} + lambda_balance = router_cfg.get("lambda_balance", 0.0) + lambda_proc = router_cfg.get("lambda_proc", 0.0) + lambda_entropy = router_cfg.get("lambda_entropy", 0.0) + gumbel_tau_start = router_cfg.get("gumbel_tau_start", 1.0) + gumbel_tau_end = router_cfg.get("gumbel_tau_end", 0.1) + + if generator == "wgan": + critic = critics.get(name) + assert critic is not None, ( + f"{name}_model.generator='wgan' requires a critic (see " + "giant.model.network.build_critics)" + ) + wgan_cfg = stage_cfg.get("wgan", {}) + trainers[name] = WGANStageTrainer( + name=name, + model=model, + critic=critic, + is_stage2=is_stage2, + lambda_weight=lambda_weight, + n_sec_lambda=n_sec_lambda, + n_critic=wgan_cfg.get("n_critic", 5), + gp_weight=wgan_cfg.get("gp_weight", 10.0), + lr=t["lr"], + critic_lr=wgan_cfg.get("critic_lr", 0.0), + ema_decay=t.get("ema_decay", 0.9999), + warmup_epochs=t["warmup_epochs"], + epochs=t["epochs"], + steps_per_epoch=max(total_train_batches, 1), + device=device, + ) + else: + ddpm_n_steps = stage_cfg.get("ddpm", {}).get("n_steps", 1000) + trainers[name] = FlowDDPMStageTrainer( + name=name, + model=model, + is_stage2=is_stage2, + generator=generator, + lambda_weight=lambda_weight, + n_sec_lambda=n_sec_lambda, + lambda_balance=lambda_balance, + lambda_proc=lambda_proc, + lambda_entropy=lambda_entropy, + gumbel_tau_start=gumbel_tau_start, + gumbel_tau_end=gumbel_tau_end, + lr=t["lr"], + weight_decay=t.get("weight_decay", 0.01), + ema_decay=t.get("ema_decay", 0.9999), + warmup_epochs=t["warmup_epochs"], + epochs=t["epochs"], + steps_per_epoch=max(total_train_batches, 1), + ddpm_n_steps=ddpm_n_steps, + device=device, + ) + return trainers + + +def _metrics_fields(trainers: dict[str, StageTrainer]) -> list[str]: + fields = ["epoch"] + for name, trainer in trainers.items(): + is_wgan = isinstance(trainer, WGANStageTrainer) + if is_wgan: + fields += [ + f"{name}_train_d_loss", + f"{name}_train_g_loss", + f"{name}_train_wasserstein", + f"{name}_train_gp_loss", + f"{name}_train_loss_nsec", + f"{name}_train_nsec_acc", + f"{name}_train_grad_norm_d", + f"{name}_train_grad_norm_g", + ] + else: + fields += [ + f"{name}_train_loss", + f"{name}_train_loss_gen", + f"{name}_train_loss_nsec", + f"{name}_train_loss_balance", + f"{name}_train_loss_proc", + f"{name}_train_loss_entropy", + f"{name}_train_nsec_acc", + f"{name}_train_grad_norm", + f"{name}_val_loss", + f"{name}_val_loss_gen", + f"{name}_val_loss_nsec", + f"{name}_val_nsec_acc", + ] + if trainer.router is not None: + fields += [ + f"{name}_router_entropy", + f"{name}_router_util_min", + f"{name}_router_util_max", + f"{name}_router_util_std", + ] + fields += [f"{name}_lr"] + if is_wgan: + fields += [f"{name}_critic_lr"] + fields += [ + "val_loss", + "val_marginal_kl", + "grad_norm", + "gpu_mem_mb", + "samples_per_sec", + "is_best", + "epoch_time_s", + ] + return fields + + +def _wandb_run_config(cfg: dict, model_config: dict | None, param_counts: dict) -> dict: return { - "d_loss": d_loss.detach(), - "g_loss": (g1 + lambda_s2 * g2).detach(), - "wasserstein_estimate": wasserstein_estimate, - "gp_loss": (gp1 + lambda_s2 * gp2).detach(), - "l_nsec": l_nsec.detach(), - "nsec_acc": nsec_acc.detach(), - "did_g_step": did_g_step, - "grad_norm": grad_norm_d.item() + grad_norm_g.item(), - "grad_norm_d": grad_norm_d.item(), - "grad_norm_g": grad_norm_g.item(), + "train": cfg["train"], + "conditioning": cfg["conditioning"], + "stage1_model": cfg["stage1_model"], + "stage2_model": cfg["stage2_model"], + "model_config": model_config or {}, + "param_counts": param_counts, } +@torch.no_grad() +def _router_gate_stats(router, cond_cont, cond_cat): + return router.gate_stats(cond_cont, cond_cat) + + def train( - stage1_model: torch.nn.Module, - sec_decoder: torch.nn.Module, + cfg: dict, + models: dict[str, torch.nn.Module | None], + critics: dict[str, torch.nn.Module | None], train_loader: DataLoader, val_loader: DataLoader, - mode: str, - epochs: int, - lr: float, - warmup_epochs: int, device: torch.device, out_dir: str | Path, - weight_decay: float = 0.01, - ema_decay: float = 0.9999, - lambda_nsec: float = 0.1, - lambda_s2: float = 1.0, - lambda_balance: float = 0.0, - lambda_proc: float = 0.0, - lambda_entropy: float = 0.0, - gumbel_tau_start: float = 1.0, - gumbel_tau_end: float = 0.1, normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, proc_map: dict | None = None, model_config: dict | None = None, resume_path: str | Path | None = None, - validate_every: int = 0, - validate_steps: int = 10, - max_val_batches: int = 0, total_train_batches: int = 0, - critic: torch.nn.Module | None = None, - sec_critic: torch.nn.Module | None = None, - n_critic: int = 5, - gp_weight: float = 10.0, - critic_lr: float | None = None, use_wandb: bool = False, wandb_project: str = "giant", wandb_run_name: str = "", wandb_log_every: int = 50, ) -> None: + """Train whichever of stage1/stage2 are active, each through its own + `StageTrainer` (design doc §7). `models`/`critics` are the dicts + `giant.model.network.build_models`/`build_critics` return — a `None` + entry means that stage is `active = false`. + """ out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) + t = cfg["train"] + epochs = t["epochs"] + validate_every = t.get("validate_every", 0) + validate_steps = t.get("validate_steps", 10) + max_val_batches = t.get("max_val_batches", 0) - stage1_params = sum(p.numel() for p in stage1_model.parameters()) - sec_decoder_params = sum(p.numel() for p in sec_decoder.parameters()) - critic_params = ( - sum(p.numel() for p in critic.parameters()) if critic is not None else 0 - ) - sec_critic_params = ( - sum(p.numel() for p in sec_critic.parameters()) if sec_critic is not None else 0 - ) - total_params = ( - stage1_params + sec_decoder_params + critic_params + sec_critic_params - ) + trainers = _build_stage_trainers(cfg, models, critics, device, total_train_batches) + if not trainers: + raise ValueError( + "no active stage — stage1_model.active and stage2_model.active are both false" + ) + + param_counts = { + name: sum(p.numel() for p in tr.model.parameters()) + for name, tr in trainers.items() + } + param_counts["total"] = sum(param_counts.values()) wandb_run = None if use_wandb: @@ -472,138 +808,32 @@ def train( "train.wandb = true (--wandb) requires the 'wandb' package — " "install it via `uv sync --extra wandb`" ) from exc - # `id` is derived from out_dir so resuming a run (--resume) reattaches - # to the same wandb run instead of starting a new one. wandb_run = wandb.init( project=wandb_project, name=wandb_run_name or out_dir.name, id=out_dir.name, resume="allow", - config=_wandb_run_config( - mode=mode, - epochs=epochs, - lr=lr, - warmup_epochs=warmup_epochs, - weight_decay=weight_decay, - ema_decay=ema_decay, - lambda_nsec=lambda_nsec, - lambda_s2=lambda_s2, - lambda_balance=lambda_balance, - lambda_proc=lambda_proc, - lambda_entropy=lambda_entropy, - gumbel_tau_start=gumbel_tau_start, - gumbel_tau_end=gumbel_tau_end, - n_critic=n_critic, - gp_weight=gp_weight, - model_config=model_config, - stage1_params=stage1_params, - sec_decoder_params=sec_decoder_params, - critic_params=critic_params, - sec_critic_params=sec_critic_params, - total_params=total_params, - ), + config=_wandb_run_config(cfg, model_config, param_counts), ) - stage1_model = stage1_model.to(device) - sec_decoder = sec_decoder.to(device) - if mode == "wgan": - assert critic is not None and sec_critic is not None, ( - "mode='wgan' requires critic/sec_critic (see giant.model.network.build_critics)" - ) - critic = critic.to(device) - sec_critic = sec_critic.to(device) - - # MoE routing trunk (RoutedDenoisingMLP/RoutedSecondaryDecoder) is - # optional and orthogonal to `mode` — both stages carry a `.router` - # when enabled. Each router is an independent instance (their - # `n_experts` need not match), used both for the batch-level gate - # entropy snapshot below and the val-level gate stats further down. - has_router = hasattr(stage1_model, "router") and hasattr(sec_decoder, "router") - - # Flow-matching/diffusion models sample noticeably better from an EMA of - # the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed - # sinusoidal-embedding freqs, or non-learned router centers) never change - # after this initial copy, so only parameters need the running average. - ema_stage1_model: torch.nn.Module | None = None - ema_sec_decoder: torch.nn.Module | None = None - if ema_decay > 0: - ema_stage1_model = copy.deepcopy(stage1_model).eval() - ema_sec_decoder = copy.deepcopy(sec_decoder).eval() - for p in ema_stage1_model.parameters(): - p.requires_grad_(False) - for p in ema_sec_decoder.parameters(): - p.requires_grad_(False) - - all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters()) - all_params_d: list = [] - optimizer_d: optim.Optimizer | None = None - if mode == "wgan": - assert critic is not None and sec_critic is not None - # Standard WGAN-GP recipe (Gulrajani et al. 2017): Adam with - # beta1=0 (momentum destabilizes critic training) and no weight - # decay, rather than the AdamW(weight_decay=...) used for flow/ddpm. - optimizer = optim.Adam(all_params, lr=lr, betas=(0.0, 0.9)) - all_params_d = list(critic.parameters()) + list(sec_critic.parameters()) - optimizer_d = optim.Adam( - all_params_d, - lr=critic_lr if critic_lr is not None else lr, - betas=(0.0, 0.9), - ) - else: - optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay) - - # Warmup/decay in units of optimizer steps rather than epochs: at large - # dataset sizes a single epoch can be tens of thousands of steps, and an - # epoch-granularity schedule would leave warmup/cosine decay unable to - # move within it. Requires an accurate `total_train_batches` (steps per - # epoch); the only caller, run_train_job, always supplies one. - # - # In wgan mode, `lr_sched.step()`/EMA only fire on generator steps (see - # the per-batch loop below) — 1 in every `n_critic` batches — so the - # schedule's own step-counting must be in those same units, or warmup - # would never finish and cosine decay would barely move. - steps_per_epoch = max(total_train_batches, 1) - if mode == "wgan": - # Generator steps fire every n_critic-th batch (did_g_step = - # step_count % n_critic == 0 in _wgan_train_step), not n_critic + 1. - steps_per_epoch = max(total_train_batches // n_critic, 1) - warmup_steps = warmup_epochs * steps_per_epoch - total_steps = max(epochs * steps_per_epoch, 1) - - def _lr_lambda(step: int) -> float: - if warmup_steps > 0 and step < warmup_steps: - return (step + 1) / warmup_steps - t = step - warmup_steps - T = max(total_steps - warmup_steps, 1) - return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T)) - - lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda) - - ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None - def _build_checkpoint(epoch: int, global_step: int, best_val_loss: float) -> dict: ckpt: dict = { - "model": stage1_model.state_dict(), - "sec_decoder": sec_decoder.state_dict(), - "optimizer": optimizer.state_dict(), - "lr_sched": lr_sched.state_dict(), "epoch": epoch, "best_val_loss": best_val_loss, "global_step": global_step, } - if mode == "wgan": - assert ( - critic is not None - and sec_critic is not None - and optimizer_d is not None - ) - ckpt["critic"] = critic.state_dict() - ckpt["sec_critic"] = sec_critic.state_dict() - ckpt["optimizer_d"] = optimizer_d.state_dict() - if ema_decay > 0: - assert ema_stage1_model is not None and ema_sec_decoder is not None - ckpt["model_ema"] = ema_stage1_model.state_dict() - ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict() + for name, tr in trainers.items(): + sd = tr.state_dict() + key = "model" if name == "stage1" else "sec_decoder" + ckpt[key] = sd["model"] + if "model_ema" in sd: + ckpt[f"{key}_ema"] = sd["model_ema"] + if "critic" in sd: + ckpt["critic" if name == "stage1" else "sec_critic"] = sd["critic"] + ckpt[f"optimizer_{name}"] = sd["optimizer"] + ckpt[f"lr_sched_{name}"] = sd["lr_sched"] + if "optimizer_d" in sd: + ckpt[f"optimizer_d_{name}"] = sd["optimizer_d"] if normalizer_dict is not None: ckpt["normalizer"] = normalizer_dict if pdg_map is not None: @@ -616,52 +846,37 @@ def train( ckpt["model_config"] = model_config return ckpt + def _load_checkpoint(ckpt: dict) -> None: + for name, tr in trainers.items(): + key = "model" if name == "stage1" else "sec_decoder" + sd = { + "model": ckpt[key], + "optimizer": ckpt[f"optimizer_{name}"], + "lr_sched": ckpt[f"lr_sched_{name}"], + } + ema_key = f"{key}_ema" + if ema_key in ckpt: + sd["model_ema"] = ckpt[ema_key] + crit_key = "critic" if name == "stage1" else "sec_critic" + if crit_key in ckpt: + sd["critic"] = ckpt[crit_key] + sd["optimizer_d"] = ckpt[f"optimizer_d_{name}"] + tr.load_state_dict(sd) + if isinstance(tr, WGANStageTrainer): + wgan_cfg = cfg[f"{name}_model"].get("wgan", {}) + tr.resume_lr(t["lr"], wgan_cfg.get("critic_lr", 0.0)) + else: + tr.resume_lr(t["lr"]) + start_epoch = 1 best_val_loss = float("inf") resumed_global_step = 0 if resume_path is not None: ckpt = torch.load(resume_path, map_location=device, weights_only=False) - stage1_model.load_state_dict(ckpt["model"]) - sec_decoder.load_state_dict(ckpt["sec_decoder"]) - if ema_decay > 0: - assert ema_stage1_model is not None and ema_sec_decoder is not None - ema_stage1_model.load_state_dict(ckpt.get("model_ema", ckpt["model"])) - ema_sec_decoder.load_state_dict( - ckpt.get("sec_decoder_ema", ckpt["sec_decoder"]) - ) - if mode == "wgan": - assert ( - critic is not None - and sec_critic is not None - and optimizer_d is not None - ) - critic.load_state_dict(ckpt["critic"]) - sec_critic.load_state_dict(ckpt["sec_critic"]) - optimizer_d.load_state_dict(ckpt["optimizer_d"]) - # Mirrors the `lr` fixup below for the generator optimizer: - # optimizer_d.load_state_dict() above restores the checkpoint's - # own critic LR, which would otherwise silently override an - # explicit `--critic-lr` passed on this resume. optimizer_d has - # no LR scheduler (unlike `optimizer`/`lr_sched`), so this is a - # flat set rather than a schedule-relative one. - resumed_critic_lr = critic_lr if critic_lr is not None else lr - for group in optimizer_d.param_groups: - group["lr"] = resumed_critic_lr - optimizer.load_state_dict(ckpt["optimizer"]) - lr_sched.load_state_dict(ckpt["lr_sched"]) + _load_checkpoint(ckpt) start_epoch = ckpt.get("epoch", 0) + 1 best_val_loss = ckpt.get("best_val_loss", float("inf")) resumed_global_step = ckpt.get("global_step", 0) - - # optimizer/lr_sched.load_state_dict() above restore the checkpoint's - # own base LR, which would otherwise silently override an explicit - # `lr` argument. Make `lr` authoritative again, applied at whatever - # point the cosine/warmup schedule has already reached. - lr_sched.base_lrs = [lr for _ in lr_sched.base_lrs] - resumed_lr = lr * _lr_lambda(lr_sched.last_epoch) - for group in optimizer.param_groups: - group["lr"] = resumed_lr - if start_epoch > epochs: print( f"checkpoint already completed epoch {start_epoch - 1} " @@ -669,51 +884,30 @@ def train( ) return + fields = _metrics_fields(trainers) metrics_path = out_dir / "metrics.csv" resuming_existing_metrics = resume_path is not None and metrics_path.exists() write_header = not resuming_existing_metrics metrics_file = open( metrics_path, "a" if resuming_existing_metrics else "w", newline="" ) - metrics_writer = csv.DictWriter(metrics_file, fieldnames=_METRICS_FIELDS) + metrics_writer = csv.DictWriter(metrics_file, fieldnames=fields) if write_header: metrics_writer.writeheader() epoch_w = len(str(epochs)) last_completed_epoch = start_epoch - 1 - # Restored from the checkpoint on --resume so wandb_run.log(..., step=...) - # keeps advancing monotonically instead of restarting at 0 mid-run (a - # reattached wandb run — see wandb.init(id=..., resume="allow") below — - # would otherwise silently drop every post-resume point). global_step = resumed_global_step with _GracefulShutdown() as shutdown: for epoch in range(start_epoch, epochs + 1): epoch_start = time.monotonic() if device.type == "cuda": torch.cuda.reset_peak_memory_stats(device) - stage1_model.train() - sec_decoder.train() - if mode == "wgan": - assert critic is not None and sec_critic is not None - critic.train() - sec_critic.train() - train_loss_sum = 0.0 - train_s1_sum = 0.0 - train_nsec_sum = 0.0 - train_s2_sum = 0.0 - train_balance_sum = 0.0 - train_proc_sum = 0.0 - train_entropy_sum = 0.0 - train_d_sum = 0.0 - train_g_sum = 0.0 - train_wasserstein_sum = 0.0 - train_gp_sum = 0.0 - train_nsec_acc_sum = 0.0 - train_grad_norm_d_sum = 0.0 - train_grad_norm_g_sum = 0.0 + for tr in trainers.values(): + tr.train_mode() + + train_sums: dict[str, dict[str, float]] = {n: {} for n in trainers} train_n = 0 - train_batches = 0 - grad_norm_sum = 0.0 ema_loss = 0.0 ema_grad_norm = 0.0 bar = tqdm( @@ -725,108 +919,33 @@ def train( dynamic_ncols=True, ) for batch in bar: - if has_router: - gumbel_tau = _gumbel_tau( - global_step, total_steps, gumbel_tau_start, gumbel_tau_end - ) - stage1_model.router.gumbel_tau = gumbel_tau - sec_decoder.router.gumbel_tau = gumbel_tau - - if mode == "wgan": - assert ( - critic is not None - and sec_critic is not None - and optimizer_d is not None - ) - stats = _wgan_train_step( - stage1_model, - sec_decoder, - critic, - sec_critic, - batch, - device, - optimizer, - optimizer_d, - all_params, - all_params_d, - global_step, - n_critic, - gp_weight, - lambda_nsec, - lambda_s2, - ) - if stats["did_g_step"]: - lr_sched.step() - if ema_decay > 0: - assert ( - ema_stage1_model is not None - and ema_sec_decoder is not None - ) - _update_ema(ema_stage1_model, stage1_model, ema_decay) - _update_ema(ema_sec_decoder, sec_decoder, ema_decay) - - B = batch[0].size(0) - batch_loss = stats["d_loss"].item() + stats["g_loss"].item() - batch_grad_norm = stats["grad_norm"] - train_loss_sum += batch_loss * B - train_nsec_sum += stats["l_nsec"].item() * B - train_d_sum += stats["d_loss"].item() * B - train_g_sum += stats["g_loss"].item() * B - train_wasserstein_sum += stats["wasserstein_estimate"].item() * B - train_gp_sum += stats["gp_loss"].item() * B - train_nsec_acc_sum += stats["nsec_acc"].item() * B - train_grad_norm_d_sum += stats["grad_norm_d"] * B - train_grad_norm_g_sum += stats["grad_norm_g"] * B - else: - loss, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc = ( - _compute_losses( - stage1_model, - sec_decoder, - batch, - mode, - ddpm_schedule, - device, - lambda_nsec, - lambda_s2, - lambda_balance, - lambda_proc, - lambda_entropy, - ) - ) - optimizer.zero_grad() - loss.backward() - grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0) - optimizer.step() - lr_sched.step() - if ema_decay > 0: - assert ( - ema_stage1_model is not None and ema_sec_decoder is not None - ) - _update_ema(ema_stage1_model, stage1_model, ema_decay) - _update_ema(ema_sec_decoder, sec_decoder, ema_decay) - - B = batch[0].size(0) - batch_loss = loss.item() - batch_grad_norm = grad_norm.item() - train_loss_sum += batch_loss * B - train_s1_sum += l_s1.item() * B - train_nsec_sum += l_nsec.item() * B - train_s2_sum += l_s2.item() * B - train_balance_sum += l_balance.item() * B - train_proc_sum += l_proc.item() * B - train_entropy_sum += l_entropy.item() * B - train_nsec_acc_sum += nsec_acc.item() * B + batch_loss_total = 0.0 + batch_grad_norm_total = 0.0 + B = batch[0].size(0) + for name, tr in trainers.items(): + stats = tr.step(batch, device, global_step) + sums = train_sums[name] + for k, v in stats.items(): + if isinstance(v, bool): + continue + sums[k] = sums.get(k, 0.0) + v * B + if isinstance(tr, WGANStageTrainer): + batch_loss_total += stats["d_loss"] + stats["g_loss"] + batch_grad_norm_total += stats["grad_norm"] + else: + batch_loss_total += stats["loss"] + batch_grad_norm_total += stats["grad_norm"] train_n += B - train_batches += 1 - grad_norm_sum += batch_grad_norm ema_loss = ( - batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss + batch_loss_total + if train_n == B + else 0.95 * ema_loss + 0.05 * batch_loss_total ) ema_grad_norm = ( - batch_grad_norm - if train_batches == 1 - else 0.95 * ema_grad_norm + 0.05 * batch_grad_norm + batch_grad_norm_total + if train_n == B + else 0.95 * ema_grad_norm + 0.05 * batch_grad_norm_total ) bar.set_postfix_str( f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False @@ -840,40 +959,22 @@ def train( ): log_payload = { "batch/epoch": epoch, - "batch/loss": batch_loss, + "batch/loss": batch_loss_total, "batch/loss_ema": ema_loss, - "batch/grad_norm": batch_grad_norm, - "batch/lr": optimizer.param_groups[0]["lr"], - "batch/critic_lr": ( - optimizer_d.param_groups[0]["lr"] - if optimizer_d is not None - else 0.0 - ), + "batch/grad_norm": batch_grad_norm_total, } - if has_router: - # Cheap re-use of the batch already in hand — no - # extra data loading, just a small forward through - # each router's own gate function. Only entropy is - # logged at this granularity (not per-expert - # utilization): a single batch's importance sum is - # too noisy as a "global share" estimate, whereas - # the val-loop aggregate (below) sums over the - # whole val set for that. Batch-level entropy alone - # is still enough to see a router collapsing in - # real time, mid-epoch, rather than only at the - # next validation pass. - with torch.no_grad(): - cond_cont_b = batch[0].to(device) - cond_cat_b = batch[1].to(device) - s1_entropy, _ = stage1_model.router.gate_stats( - cond_cont_b, cond_cat_b - ) - s2_entropy, _ = sec_decoder.router.gate_stats( - cond_cont_b, cond_cat_b - ) - log_payload["batch/router_s1_entropy"] = s1_entropy.item() - log_payload["batch/router_s2_entropy"] = s2_entropy.item() - log_payload["batch/gumbel_tau"] = gumbel_tau + for name, tr in trainers.items(): + log_payload[f"batch/{name}_lr"] = tr.optimizer.param_groups[0][ + "lr" + ] + if tr.router is not None: + with torch.no_grad(): + cond_cont_b = batch[0].to(device) + cond_cat_b = batch[1].to(device) + entropy, _ = _router_gate_stats( + tr.router, cond_cont_b, cond_cat_b + ) + log_payload[f"batch/{name}_router_entropy"] = entropy.item() wandb_run.log(log_payload, step=global_step) if shutdown.requested: @@ -881,12 +982,6 @@ def train( bar.close() if shutdown.requested: - # Epoch was interrupted mid-loop, so there's no val_loss to - # weigh a "best" checkpoint against — save the in-progress - # weights as last.pt only, under the last *fully completed* - # epoch number so --resume restarts this epoch from scratch - # rather than skipping it (weights/optimizer state are still - # kept, so those partial-epoch batches aren't wasted work). ckpt = _build_checkpoint(epoch - 1, global_step, best_val_loss) torch.save(ckpt, out_dir / "last.pt") last_completed_epoch = epoch - 1 @@ -897,165 +992,112 @@ def train( ) break - train_loss = train_loss_sum / max(train_n, 1) - train_grad_norm = grad_norm_sum / max(train_batches, 1) - train_nsec_acc = train_nsec_acc_sum / max(train_n, 1) - train_grad_norm_d = train_grad_norm_d_sum / max(train_n, 1) - train_grad_norm_g = train_grad_norm_g_sum / max(train_n, 1) - current_lr = optimizer.param_groups[0]["lr"] - critic_lr_value = ( - optimizer_d.param_groups[0]["lr"] if optimizer_d is not None else 0.0 - ) + for tr in trainers.values(): + tr.eval_mode() - stage1_model.eval() - sec_decoder.eval() - if mode == "wgan": - assert critic is not None and sec_critic is not None - critic.eval() - sec_critic.eval() - - val_marginal_kl = float("nan") - if mode == "wgan": - # WGANGenerator.forward(z, cond_cont, cond_cat) has no - # diffusion/flow `t` argument, so the usual _compute_losses - # val loop below (which calls flow_matching_loss -> - # stage1_model(x_t, t, ...)) doesn't apply — and a WGAN - # critic loss isn't a monotone quality signal fit for - # best-checkpoint selection anyway. Select on marginal KL - # against the EMA generators instead (matches what - # predict/rollout sample from by default, --weights ema). - eval_stage1 = ( - ema_stage1_model if ema_stage1_model is not None else stage1_model - ) - eval_sec_decoder = ( - ema_sec_decoder if ema_sec_decoder is not None else sec_decoder - ) - marginal_result = validate_marginals( - eval_stage1, - val_loader, - mode=mode, - device=device, - sec_decoder=eval_sec_decoder, - ) - val_marginal_kl = float(np.mean(marginal_result["kl_divergence"])) - val_loss = val_marginal_kl - val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = ( - val_proc_sum - ) = val_entropy_sum = val_nsec_acc_sum = 0.0 - val_n = 1 - val_nsec_acc = 0.0 - router_s1_entropy = router_s2_entropy = 0.0 - router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0 - router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0 - else: - val_loss_sum = 0.0 - val_s1_sum = 0.0 - val_nsec_sum = 0.0 - val_s2_sum = 0.0 - val_balance_sum = 0.0 - val_proc_sum = 0.0 - val_entropy_sum = 0.0 - val_nsec_acc_sum = 0.0 - val_n = 0 - if has_router: - n_experts_s1 = stage1_model.router.n_experts - n_experts_s2 = sec_decoder.router.n_experts - val_router_s1_entropy_sum = 0.0 - val_router_s2_entropy_sum = 0.0 - val_router_s1_importance_sum = torch.zeros( - n_experts_s1, device=device - ) - val_router_s2_importance_sum = torch.zeros( - n_experts_s2, device=device - ) + # --- per-stage validation --- + val_sums: dict[str, dict[str, float]] = {n: {} for n in trainers} + val_router_sums: dict[str, dict] = {} + val_n = 0 + non_adversarial = { + n: tr + for n, tr in trainers.items() + if not isinstance(tr, WGANStageTrainer) + } + if non_adversarial: with torch.no_grad(): for val_batch_idx, batch in enumerate(val_loader): if max_val_batches > 0 and val_batch_idx >= max_val_batches: break - ( - loss, - l_s1, - l_nsec, - l_s2, - l_balance, - l_proc, - l_entropy, - nsec_acc, - ) = _compute_losses( - stage1_model, - sec_decoder, - batch, - mode, - ddpm_schedule, - device, - lambda_nsec, - lambda_s2, - lambda_balance, - lambda_proc, - lambda_entropy, - ) B = batch[0].size(0) - val_loss_sum += loss.item() * B - val_s1_sum += l_s1.item() * B - val_nsec_sum += l_nsec.item() * B - val_s2_sum += l_s2.item() * B - val_balance_sum += l_balance.item() * B - val_proc_sum += l_proc.item() * B - val_entropy_sum += l_entropy.item() * B - val_nsec_acc_sum += nsec_acc.item() * B - if has_router: - cond_cont = batch[0].to(device) - cond_cat = batch[1].to(device) - s1_entropy, s1_importance = stage1_model.router.gate_stats( - cond_cont, cond_cat - ) - s2_entropy, s2_importance = sec_decoder.router.gate_stats( - cond_cont, cond_cat - ) - val_router_s1_entropy_sum += s1_entropy.item() * B - val_router_s2_entropy_sum += s2_entropy.item() * B - val_router_s1_importance_sum += s1_importance - val_router_s2_importance_sum += s2_importance + for name, tr in non_adversarial.items(): + stats = tr.val_loss(batch, device) + sums = val_sums[name] + for k, v in stats.items(): + sums[k] = sums.get(k, 0.0) + v * B + if tr.router is not None: + cond_cont = batch[0].to(device) + cond_cat = batch[1].to(device) + entropy, importance = _router_gate_stats( + tr.router, cond_cont, cond_cat + ) + rs = val_router_sums.setdefault( + name, + { + "entropy": 0.0, + "importance": torch.zeros_like(importance), + "n_experts": tr.router.n_experts, + }, + ) + rs["entropy"] += entropy.item() * B + rs["importance"] += importance val_n += B - val_loss = val_loss_sum / max(val_n, 1) - val_nsec_acc = val_nsec_acc_sum / max(val_n, 1) - if has_router: - router_s1_entropy = val_router_s1_entropy_sum / max(val_n, 1) - router_s2_entropy = val_router_s2_entropy_sum / max(val_n, 1) - s1_util = val_router_s1_importance_sum / ( - val_router_s1_importance_sum.sum().clamp_min(1e-8) + val_loss_per_stage: dict[str, float] = {} + for name, sums in val_sums.items(): + n = max(val_n, 1) + val_loss_per_stage[name] = sums.get("loss", 0.0) / n + + val_marginal_kl = float("nan") + wgan_names = [ + n for n, tr in trainers.items() if isinstance(tr, WGANStageTrainer) + ] + if wgan_names: + # No monotone per-batch WGAN loss for best-checkpoint + # selection — try the real marginal-KL signal (see + # _try_validate_marginals), falling back to the epoch's own + # Wasserstein-distance magnitude, averaged over active wgan + # stages, if that's unavailable. + stage1_tr = trainers.get("stage1") + kl_result = None + if stage1_tr is not None: + kl_result = _try_validate_marginals( + stage1_tr, + val_loader, + device, + mode=cfg["stage1_model"]["generator"], + sec_decoder=trainers["stage2"].sampling_model() + if "stage2" in trainers + else None, ) - s2_util = val_router_s2_importance_sum / ( - val_router_s2_importance_sum.sum().clamp_min(1e-8) - ) - router_s1_util_min = s1_util.min().item() - router_s1_util_max = s1_util.max().item() - router_s1_util_std = ( - s1_util.std().item() if n_experts_s1 > 1 else 0.0 - ) - router_s2_util_min = s2_util.min().item() - router_s2_util_max = s2_util.max().item() - router_s2_util_std = ( - s2_util.std().item() if n_experts_s2 > 1 else 0.0 + if kl_result is not None: + val_marginal_kl = float(np.mean(kl_result["kl_divergence"])) + val_loss_per_stage.update( + { + n: val_marginal_kl + for n in wgan_names + if n not in val_loss_per_stage + } ) else: - router_s1_entropy = router_s2_entropy = 0.0 - router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0 - router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0 - - if validate_every > 0 and epoch % validate_every == 0: - print(f"[epoch {epoch}] marginal validation:") - marginal_result = validate_marginals( - stage1_model, + wasserstein_proxy = sum( + train_sums[n].get("wasserstein_estimate", 0.0) + for n in wgan_names + ) / max(train_n, 1) + for n in wgan_names: + val_loss_per_stage[n] = abs(wasserstein_proxy) + elif validate_every > 0 and epoch % validate_every == 0: + stage1_tr = trainers.get("stage1") + if stage1_tr is not None: + kl_result = _try_validate_marginals( + stage1_tr, val_loader, - mode=mode, - schedule=ddpm_schedule, - device=device, + device, + mode=cfg["stage1_model"]["generator"], + schedule=stage1_tr.ddpm_schedule + if isinstance(stage1_tr, FlowDDPMStageTrainer) + else None, steps=validate_steps, - sec_decoder=sec_decoder, + sec_decoder=trainers["stage2"].sampling_model() + if "stage2" in trainers + else None, ) - val_marginal_kl = float(np.mean(marginal_result["kl_divergence"])) + if kl_result is not None: + val_marginal_kl = float(np.mean(kl_result["kl_divergence"])) + + val_loss = ( + sum(val_loss_per_stage.values()) if val_loss_per_stage else float("nan") + ) epoch_time = time.monotonic() - epoch_start gpu_mem_mb = ( @@ -1066,78 +1108,129 @@ def train( is_best = val_loss < best_val_loss marker = " [best]" if is_best else "" + summary_bits = [] + for name, tr in trainers.items(): + if isinstance(tr, WGANStageTrainer): + n = max(train_n, 1) + summary_bits.append( + f"{name}[d={train_sums[name].get('d_loss', 0.0) / n:.3f} " + f"g={train_sums[name].get('g_loss', 0.0) / n:.3f}]" + ) + else: + n = max(train_n, 1) + summary_bits.append( + f"{name}[loss={train_sums[name].get('loss', 0.0) / n:.3f}]" + ) print( - f"epoch {epoch:{epoch_w}d}/{epochs}" - f" train {train_loss:.4f}" - f" (s1={train_s1_sum / max(train_n, 1):.3f}" - f" nsec={train_nsec_sum / max(train_n, 1):.3f}" - f" s2={train_s2_sum / max(train_n, 1):.3f}" - f" bal={train_balance_sum / max(train_n, 1):.3f}" - f" proc={train_proc_sum / max(train_n, 1):.3f}" - f" entropy={train_entropy_sum / max(train_n, 1):.3f}" - f" d={train_d_sum / max(train_n, 1):.3f}" - f" g={train_g_sum / max(train_n, 1):.3f})" - f" val {val_loss:.4f}" - f" lr {current_lr:.2e} gnorm {train_grad_norm:.3f}" - f" {epoch_time:.1f}s{marker}" + f"epoch {epoch:{epoch_w}d}/{epochs} " + + " ".join(summary_bits) + + f" val {val_loss:.4f} {epoch_time:.1f}s{marker}" ) - metrics_row = { - "epoch": epoch, - "train_loss": train_loss, - "train_loss_s1": train_s1_sum / max(train_n, 1), - "train_loss_nsec": train_nsec_sum / max(train_n, 1), - "train_loss_s2": train_s2_sum / max(train_n, 1), - "train_loss_balance": train_balance_sum / max(train_n, 1), - "train_loss_proc": train_proc_sum / max(train_n, 1), - "train_loss_entropy": train_entropy_sum / max(train_n, 1), - "train_nsec_acc": train_nsec_acc, - "d_loss": train_d_sum / max(train_n, 1), - "g_loss": train_g_sum / max(train_n, 1), - "wasserstein_estimate": train_wasserstein_sum / max(train_n, 1), - "gp_loss": train_gp_sum / max(train_n, 1), - "val_loss": val_loss, - "val_loss_s1": val_s1_sum / max(val_n, 1), - "val_loss_nsec": val_nsec_sum / max(val_n, 1), - "val_loss_s2": val_s2_sum / max(val_n, 1), - "val_loss_balance": val_balance_sum / max(val_n, 1), - "val_loss_proc": val_proc_sum / max(val_n, 1), - "val_loss_entropy": val_entropy_sum / max(val_n, 1), - "val_nsec_acc": val_nsec_acc, - "val_marginal_kl": val_marginal_kl, - "router_s1_entropy": router_s1_entropy, - "router_s1_util_min": router_s1_util_min, - "router_s1_util_max": router_s1_util_max, - "router_s1_util_std": router_s1_util_std, - "router_s2_entropy": router_s2_entropy, - "router_s2_util_min": router_s2_util_min, - "router_s2_util_max": router_s2_util_max, - "router_s2_util_std": router_s2_util_std, - "lr": current_lr, - "critic_lr": critic_lr_value, - "grad_norm": train_grad_norm, - "grad_norm_d": train_grad_norm_d, - "grad_norm_g": train_grad_norm_g, - "gpu_mem_mb": gpu_mem_mb, - "samples_per_sec": train_n / max(epoch_time, 1e-8), - "is_best": int(is_best), - "epoch_time_s": epoch_time, - } + + metrics_row: dict = {"epoch": epoch} + grad_norm_total = 0.0 + for name, tr in trainers.items(): + n_train = max(train_n, 1) + sums = train_sums[name] + if isinstance(tr, WGANStageTrainer): + metrics_row[f"{name}_train_d_loss"] = ( + sums.get("d_loss", 0.0) / n_train + ) + metrics_row[f"{name}_train_g_loss"] = ( + sums.get("g_loss", 0.0) / n_train + ) + metrics_row[f"{name}_train_wasserstein"] = ( + sums.get("wasserstein_estimate", 0.0) / n_train + ) + metrics_row[f"{name}_train_gp_loss"] = ( + sums.get("gp_loss", 0.0) / n_train + ) + metrics_row[f"{name}_train_loss_nsec"] = ( + sums.get("loss_nsec", 0.0) / n_train + ) + metrics_row[f"{name}_train_nsec_acc"] = ( + sums.get("nsec_acc", 0.0) / n_train + ) + metrics_row[f"{name}_train_grad_norm_d"] = ( + sums.get("grad_norm_d", 0.0) / n_train + ) + metrics_row[f"{name}_train_grad_norm_g"] = ( + sums.get("grad_norm_g", 0.0) / n_train + ) + grad_norm_total += sums.get("grad_norm", 0.0) / n_train + metrics_row[f"{name}_critic_lr"] = tr.optimizer_d.param_groups[0][ + "lr" + ] + else: + n_val = max(val_n, 1) + v = val_sums.get(name, {}) + metrics_row[f"{name}_train_loss"] = sums.get("loss", 0.0) / n_train + metrics_row[f"{name}_train_loss_gen"] = ( + sums.get("loss_gen", 0.0) / n_train + ) + metrics_row[f"{name}_train_loss_nsec"] = ( + sums.get("loss_nsec", 0.0) / n_train + ) + metrics_row[f"{name}_train_loss_balance"] = ( + sums.get("loss_balance", 0.0) / n_train + ) + metrics_row[f"{name}_train_loss_proc"] = ( + sums.get("loss_proc", 0.0) / n_train + ) + metrics_row[f"{name}_train_loss_entropy"] = ( + sums.get("loss_entropy", 0.0) / n_train + ) + metrics_row[f"{name}_train_nsec_acc"] = ( + sums.get("nsec_acc", 0.0) / n_train + ) + metrics_row[f"{name}_train_grad_norm"] = ( + sums.get("grad_norm", 0.0) / n_train + ) + metrics_row[f"{name}_val_loss"] = v.get("loss", 0.0) / n_val + metrics_row[f"{name}_val_loss_gen"] = v.get("loss_gen", 0.0) / n_val + metrics_row[f"{name}_val_loss_nsec"] = ( + v.get("loss_nsec", 0.0) / n_val + ) + metrics_row[f"{name}_val_nsec_acc"] = v.get("nsec_acc", 0.0) / n_val + grad_norm_total += sums.get("grad_norm", 0.0) / n_train + if tr.router is not None: + rs = val_router_sums.get(name) + if rs is not None: + n_experts = rs["n_experts"] + util = rs["importance"] / rs["importance"].sum().clamp_min(1e-8) + metrics_row[f"{name}_router_entropy"] = rs["entropy"] / max( + val_n, 1 + ) + metrics_row[f"{name}_router_util_min"] = util.min().item() + metrics_row[f"{name}_router_util_max"] = util.max().item() + metrics_row[f"{name}_router_util_std"] = ( + util.std().item() if n_experts > 1 else 0.0 + ) + else: + metrics_row[f"{name}_router_entropy"] = 0.0 + metrics_row[f"{name}_router_util_min"] = 0.0 + metrics_row[f"{name}_router_util_max"] = 0.0 + metrics_row[f"{name}_router_util_std"] = 0.0 + metrics_row[f"{name}_lr"] = tr.optimizer.param_groups[0]["lr"] + + metrics_row["val_loss"] = val_loss + metrics_row["val_marginal_kl"] = val_marginal_kl + metrics_row["grad_norm"] = grad_norm_total + metrics_row["gpu_mem_mb"] = gpu_mem_mb + metrics_row["samples_per_sec"] = train_n / max(epoch_time, 1e-8) + metrics_row["is_best"] = int(is_best) + metrics_row["epoch_time_s"] = epoch_time + metrics_writer.writerow(metrics_row) metrics_file.flush() if wandb_run is not None: - # Shares the same monotonic step axis as the per-batch - # `batch/*` logs above (global_step) rather than `epoch`, - # since a wandb run's `step` argument across `log()` calls - # must never decrease. wandb_run.log(metrics_row, step=global_step) ckpt = _build_checkpoint(epoch, global_step, best_val_loss) - - if val_loss < best_val_loss: + if is_best: best_val_loss = val_loss ckpt["best_val_loss"] = best_val_loss torch.save(ckpt, out_dir / "best.pt") - torch.save(ckpt, out_dir / "last.pt") last_completed_epoch = epoch diff --git a/scripts/warm_setup_cache.py b/scripts/warm_setup_cache.py index 3eeebb3..030529b 100644 --- a/scripts/warm_setup_cache.py +++ b/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, diff --git a/tests/test_cli_new_run.py b/tests/test_cli_new_run.py index 54788c4..c6ef14f 100644 --- a/tests/test_cli_new_run.py +++ b/tests/test_cli_new_run.py @@ -37,13 +37,14 @@ def test_writes_config_with_overrides_applied(tmp_path: Path): with open(config_path, "rb") as f: cfg = tomllib.load(f) - assert cfg["train"]["mode"] == "ddpm" + assert cfg["stage1_model"]["generator"] == "ddpm" + assert cfg["stage2_model"]["generator"] == "ddpm" assert cfg["train"]["lr"] == 0.0005 - assert cfg["model"]["hidden_dim"] == 128 - assert cfg["model"]["n_blocks"] == 4 + assert cfg["stage1_model"]["hidden_dim"] == 128 + assert cfg["stage1_model"]["n_res_blocks"] == 4 # untouched defaults still present assert cfg["train"]["epochs"] == 100 - assert "router" in cfg["model"] + assert "router" in cfg["stage1_model"] assert str(out_dir) in result.output assert "" in result.output diff --git a/tests/test_config.py b/tests/test_config.py index eee0f18..653fca1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -576,11 +576,29 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning(): **{ "stage2_model.particle_type.target": "embedding", "conditioning.particle.type": "embedding", + "conditioning.material.type": "embedding", } ) gconfig.validate_config(cfg) # must not raise +def test_validate_config_mixed_particle_material_conditioning_not_supported(): + """The data pipeline doesn't support mixed conditioning types yet, even + though ConditionEncoder itself already can (docs/v0.3.0-design.md §3.1 + vs. giant/data/transforms.py's still-single conditioning param).""" + cfg = _cfg_with( + **{ + "conditioning.particle.type": "physical", + "conditioning.material.type": "embedding", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "mixed" in str(e) or "conditioning.material.type" in str(e) + + def test_validate_config_pdg_router_incompatible_with_physical_conditioning(): cfg = _cfg_with( **{ diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 4c84c45..833b8fd 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -99,10 +99,19 @@ def _tiny_cfg(**train_overrides): "warmup_epochs": 0, "validate_every": 0, "max_val_batches": 1, + "wandb": False, } ) cfg["train"].update(train_overrides) - cfg["model"].update({"hidden_dim": 8, "n_blocks": 1, "emb_dim": 4, "dropout": 0.0}) + cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0}) + cfg["stage2_model"].update( + # decoder="autoregressive" is DEFAULT_CONFIG's default (the finished + # v0.3.0 target) but Stage2Autoregressive isn't implemented until + # design doc step 4/5 — every run must override to "one_shot" for now. + {"decoder": "one_shot", "hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0} + ) + cfg["conditioning"]["particle"]["emb_dim"] = 4 + cfg["conditioning"]["material"]["emb_dim"] = 4 return cfg diff --git a/tests/test_rollout.py b/tests/test_rollout.py index bdc0006..0e71b90 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -21,7 +21,10 @@ from giant import geometry as g # noqa: E402 # models with the pre-refactor positional convention # (model(x, t, cond_cont, cond_cat)) that no longer matches these classes' # forward signatures. Deferred to docs/v0.3.0-design.md step 6/§10. -pytestmark = pytest.mark.xfail(reason="giant/rollout.py not updated for v0.3.0 network.py yet (step 6)", strict=False) +pytestmark = pytest.mark.xfail( + reason="giant/rollout.py not updated for v0.3.0 network.py yet (step 6)", + strict=False, +) PDG_MAP = {22: 0, 11: 1, -11: 2} MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1} diff --git a/tests/test_router.py b/tests/test_router.py index 4e8b044..96b8c5b 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -841,6 +841,7 @@ def test_build_models_routed_with_composed_router(): ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None assert isinstance(stage1.trunk, RoutedTrunk) assert isinstance(stage1.trunk.router, ComposedRouter) assert len(stage1.trunk.experts) == 12 @@ -863,7 +864,9 @@ def test_build_models_routed_stage2_ties_to_stage1_router(): }, ) models = build_models(cfg) - assert models["stage1"].trunk.router is models["stage2"].trunk.router + stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None + assert stage1.trunk.router is stage2.trunk.router @pytest.mark.xfail( @@ -892,6 +895,7 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None B = 5 cond_cont, cond_cat = _cond(B, pdg=3, mat=2) stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) @@ -1021,10 +1025,11 @@ def test_routed_stage2_gradients_flow(): def test_build_models_monolith_when_router_absent(): cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2) models = build_models(cfg) - assert isinstance(models["stage1"], Stage1Model) - assert isinstance(models["stage2"], Stage2OneShot) - assert isinstance(models["stage1"].trunk, MonolithicTrunk) - assert isinstance(models["stage2"].trunk, MonolithicTrunk) + stage1, stage2 = models["stage1"], models["stage2"] + assert isinstance(stage1, Stage1Model) + assert isinstance(stage2, Stage2OneShot) + assert isinstance(stage1.trunk, MonolithicTrunk) + assert isinstance(stage2.trunk, MonolithicTrunk) def test_build_models_monolith_when_router_disabled(): @@ -1034,8 +1039,10 @@ def test_build_models_monolith_when_router_disabled(): stage1_router={"enabled": False, "type": "energy", "n_experts": 4}, ) models = build_models(cfg) - assert isinstance(models["stage1"].trunk, MonolithicTrunk) - assert isinstance(models["stage2"].trunk, MonolithicTrunk) + stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None + assert isinstance(stage1.trunk, MonolithicTrunk) + assert isinstance(stage2.trunk, MonolithicTrunk) def test_build_models_routed_when_enabled(): @@ -1061,6 +1068,7 @@ def test_build_models_routed_when_enabled(): ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None assert isinstance(stage1.trunk, RoutedTrunk) assert isinstance(stage2.trunk, RoutedTrunk) assert len(stage1.trunk.experts) == 4 @@ -1087,6 +1095,7 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow(): ) models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None B = 5 cond_cont, cond_cat = _cond(B, pdg=3, mat=2) stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2) diff --git a/tests/test_router_gating.py b/tests/test_router_gating.py index 51123d3..18be8c8 100644 --- a/tests/test_router_gating.py +++ b/tests/test_router_gating.py @@ -36,7 +36,8 @@ def _model_cfg() -> dict: def _write_checkpoint(tmp_path) -> str: cfg = _model_cfg() - stage1, _ = build_models(cfg) + stage1 = build_models(cfg)["stage1"] + assert stage1 is not None norm = Normalizer() norm.mean = np.zeros(15, dtype=np.float32) norm.std = np.ones(15, dtype=np.float32) diff --git a/tests/test_train.py b/tests/test_train.py index 627fe38..5892de7 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -1,6 +1,24 @@ -"""Tests for giant/train.py helpers.""" +"""Tests for giant/train.py.""" -from giant.train import _gumbel_tau, _wandb_run_config +import copy +import tempfile +from pathlib import Path + +import pytest +import torch + +from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM +from giant.model.network import build_critics, build_models +from giant.train import ( + FlowDDPMStageTrainer, + WGANStageTrainer, + _gumbel_tau, + _wandb_run_config, + train, +) + +PDG_VOCAB = 6 +MAT_VOCAB = 3 def test_gumbel_tau_at_step_zero_is_start(): @@ -26,72 +44,320 @@ def test_gumbel_tau_handles_zero_total_steps(): assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9 -def _base_wandb_kwargs(**overrides): - kwargs = dict( - mode="flow", - epochs=30, - lr=3e-4, - warmup_epochs=3, - weight_decay=0.01, - ema_decay=0.9999, - lambda_nsec=0.1, - lambda_s2=1.0, - lambda_balance=0.035, - lambda_proc=0.0, - lambda_entropy=0.0, - gumbel_tau_start=1.0, - gumbel_tau_end=0.1, - n_critic=5, - gp_weight=10.0, - model_config={"router": {"enabled": False}}, - stage1_params=100, - sec_decoder_params=50, - critic_params=0, - sec_critic_params=0, - total_params=150, +def test_wandb_run_config_includes_full_cfg_and_param_counts(): + cfg = { + "train": {"lr": 3e-4}, + "conditioning": {"out_dim": 128}, + "stage1_model": {"generator": "flow"}, + "stage2_model": {"generator": "wgan"}, + } + wcfg = _wandb_run_config( + cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100} ) - kwargs.update(overrides) - return kwargs - - -def test_wandb_run_config_omits_router_knobs_when_router_disabled(): - cfg = _wandb_run_config(**_base_wandb_kwargs()) - for key in ( - "lambda_balance", - "lambda_proc", - "lambda_entropy", - "gumbel_tau_start", - "gumbel_tau_end", - ): - assert key not in cfg - # still present, nested, regardless of router state - assert cfg["model"] == {"router": {"enabled": False}} - - -def test_wandb_run_config_includes_router_knobs_when_router_enabled(): - cfg = _wandb_run_config( - **_base_wandb_kwargs(model_config={"router": {"enabled": True}}) - ) - assert cfg["lambda_balance"] == 0.035 - assert cfg["lambda_proc"] == 0.0 - assert cfg["lambda_entropy"] == 0.0 - assert cfg["gumbel_tau_start"] == 1.0 - assert cfg["gumbel_tau_end"] == 0.1 - - -def test_wandb_run_config_omits_wgan_knobs_when_mode_is_not_wgan(): - cfg = _wandb_run_config(**_base_wandb_kwargs(mode="flow")) - assert "n_critic" not in cfg - assert "gp_weight" not in cfg - - -def test_wandb_run_config_includes_wgan_knobs_when_mode_is_wgan(): - cfg = _wandb_run_config(**_base_wandb_kwargs(mode="wgan")) - assert cfg["n_critic"] == 5 - assert cfg["gp_weight"] == 10.0 + assert wcfg["train"] == {"lr": 3e-4} + assert wcfg["stage1_model"] == {"generator": "flow"} + assert wcfg["stage2_model"] == {"generator": "wgan"} + assert wcfg["model_config"] == {"pdg_vocab": 3} + assert wcfg["param_counts"] == {"stage1": 100} def test_wandb_run_config_handles_missing_model_config(): - cfg = _wandb_run_config(**_base_wandb_kwargs(model_config=None)) - assert cfg["model"] == {} - assert "lambda_balance" not in cfg + cfg = {"train": {}, "conditioning": {}, "stage1_model": {}, "stage2_model": {}} + wcfg = _wandb_run_config(cfg, model_config=None, param_counts={}) + assert wcfg["model_config"] == {} + + +# --- end-to-end train() integration tests ----------------------------------- + +PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} + + +def _base_cfg(): + return { + "conditioning": { + "out_dim": 32, + "share_stages": False, + "particle": dict(PARTICLE_CFG), + "material": dict(MATERIAL_CFG), + }, + "stage1_model": { + "active": True, + "generator": "flow", + "hidden_dim": 24, + "n_res_blocks": 2, + "dropout": 0.0, + "lambda": 1.0, + "flow": {"time_dim": 16}, + "ddpm": {"time_dim": 16, "n_steps": 50}, + "wgan": { + "noise_dim": 16, + "n_critic": 2, + "gp_weight": 10.0, + "critic_lr": 0.0, + }, + "router": {"enabled": False}, + }, + "stage2_model": { + "active": True, + "decoder": "one_shot", + "generator": "wgan", + "hidden_dim": 24, + "n_res_blocks": 2, + "dropout": 0.0, + "lambda": 1.0, + "k_max": K_MAX, + "context_dim": 16, + "n_sec": {"mode": "head", "lambda": 0.1}, + "flow": {"time_dim": 16}, + "ddpm": {"time_dim": 16, "n_steps": 50}, + "wgan": { + "noise_dim": 16, + "n_critic": 2, + "gp_weight": 10.0, + "critic_lr": 0.0, + }, + "router": {"enabled": False, "tie_to_stage1": False}, + }, + "train": { + "epochs": 2, + "batch_size": 8, + "lr": 3e-4, + "weight_decay": 0.01, + "ema_decay": 0.999, + "warmup_epochs": 0, + "val_fraction": 0.1, + "max_val_batches": 0, + "num_workers": 0, + "seed": 0, + "validate_every": 0, + "validate_steps": 2, + "wandb": False, + }, + } + + +def _fake_batches(n_batches, batch_size, seed=0): + g = torch.Generator().manual_seed(seed) + batches = [] + for _ in range(n_batches): + cond_cont = torch.randn(batch_size, COND_DIM, generator=g) + cond_cat = torch.stack( + [ + torch.randint(0, PDG_VOCAB, (batch_size,), generator=g), + torch.randint(0, MAT_VOCAB, (batch_size,), generator=g), + ], + dim=1, + ) + x1 = torch.randn(batch_size, X_DIM, generator=g) + n_sec = torch.randint(0, K_MAX, (batch_size,), generator=g) + sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g) + proc_idx = torch.zeros(batch_size, dtype=torch.long) + batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx)) + return batches + + +def _model_config(cfg): + return { + "pdg_vocab": PDG_VOCAB, + "mat_vocab": MAT_VOCAB, + "conditioning": cfg["conditioning"], + "stage1_model": cfg["stage1_model"], + "stage2_model": cfg["stage2_model"], + } + + +def _run_train(cfg, out_dir, resume_path=None): + model_config = _model_config(cfg) + models = build_models(model_config) + critics = build_critics(model_config) + train_loader = _fake_batches(4, cfg["train"]["batch_size"]) + val_loader = _fake_batches(2, cfg["train"]["batch_size"], seed=1) + train( + cfg=cfg, + models=models, + critics=critics, + train_loader=train_loader, + val_loader=val_loader, + device=torch.device("cpu"), + out_dir=out_dir, + normalizer_dict={"cond": {}, "target": {}, "sec_phys": {}}, + pdg_map={"22": 0}, + mat_map={"G4_AIR": 0}, + proc_map=None, + model_config=model_config, + total_train_batches=4, + resume_path=resume_path, + ) + + +@pytest.mark.parametrize( + "label,mutate", + [ + ("both_flow", lambda cfg: None), + ("both_wgan", lambda cfg: cfg["stage1_model"].__setitem__("generator", "wgan")), + ( + "mixed_stage1_flow_stage2_wgan", + lambda cfg: None, # already the default + ), + ( + "mixed_stage1_wgan_stage2_flow", + lambda cfg: ( + cfg["stage1_model"].__setitem__("generator", "wgan"), + cfg["stage2_model"].__setitem__("generator", "flow"), + ), + ), + ("stage1_only", lambda cfg: cfg["stage2_model"].__setitem__("active", False)), + ("stage2_only", lambda cfg: cfg["stage1_model"].__setitem__("active", False)), + ( + "both_ddpm_stage1_flow_stage2", + lambda cfg: ( + cfg["stage1_model"].__setitem__("generator", "ddpm"), + cfg["stage2_model"].__setitem__("generator", "flow"), + ), + ), + ( + "routed_stage1_energy_gumbel", + lambda cfg: cfg["stage1_model"].__setitem__( + "router", + { + "enabled": True, + "type": "energy", + "n_experts": 3, + "temperature": 0.5, + "learn_centers": True, + "lambda_balance": 0.1, + "lambda_entropy": 0.01, + "gumbel": True, + "gumbel_tau_start": 1.0, + "gumbel_tau_end": 0.1, + }, + ), + ), + ], +) +def test_train_end_to_end(label, mutate): + cfg = _base_cfg() + mutate(cfg) + with tempfile.TemporaryDirectory() as tmp: + out_dir = Path(tmp) / "run" + _run_train(cfg, out_dir) + assert (out_dir / "last.pt").exists() + assert (out_dir / "metrics.csv").exists() + ckpt = torch.load(out_dir / "last.pt", weights_only=False) + if cfg["stage1_model"]["active"]: + assert "model" in ckpt + else: + assert "model" not in ckpt + if cfg["stage2_model"]["active"]: + assert "sec_decoder" in ckpt + else: + assert "sec_decoder" not in ckpt + + +def test_train_resume_continues_from_checkpoint(): + cfg = _base_cfg() + with tempfile.TemporaryDirectory() as tmp: + out_dir = Path(tmp) / "run" + _run_train(cfg, out_dir) + ckpt_before = torch.load(out_dir / "last.pt", weights_only=False) + assert ckpt_before["epoch"] == 2 + + cfg2 = copy.deepcopy(cfg) + cfg2["train"]["epochs"] = 3 + _run_train(cfg2, out_dir, resume_path=out_dir / "last.pt") + ckpt_after = torch.load(out_dir / "last.pt", weights_only=False) + assert ckpt_after["epoch"] == 3 + assert ckpt_after["global_step"] > ckpt_before["global_step"] + + +def test_train_raises_when_no_active_stage(): + cfg = _base_cfg() + cfg["stage1_model"]["active"] = False + cfg["stage2_model"]["active"] = False + model_config = _model_config(cfg) + models = build_models(model_config) + critics = build_critics(model_config) + with tempfile.TemporaryDirectory() as tmp: + with pytest.raises(ValueError, match="no active stage"): + train( + cfg=cfg, + models=models, + critics=critics, + train_loader=_fake_batches(1, 8), + val_loader=_fake_batches(1, 8), + device=torch.device("cpu"), + out_dir=Path(tmp) / "run", + model_config=model_config, + total_train_batches=1, + ) + + +def test_metrics_csv_columns_are_stage_prefixed(): + cfg = _base_cfg() + with tempfile.TemporaryDirectory() as tmp: + 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 "epoch" in header + + +def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch(): + """Regression test: on a non-generator-step batch, if this stage's model + has no n_sec_head (n_sec defaults to stage 2, decision 1), g_loss is a + graph-less zero — .backward() must not be called on it.""" + cfg = _base_cfg() + cfg["stage1_model"]["generator"] = "wgan" + model_config = _model_config(cfg) + models = build_models(model_config) + critics = build_critics(model_config) + assert models["stage1"] is not None and critics["stage1"] is not None + trainer = WGANStageTrainer( + name="stage1", + model=models["stage1"], + critic=critics["stage1"], + is_stage2=False, + lambda_weight=1.0, + n_sec_lambda=0.1, + 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"), + ) + assert trainer.model.n_sec_head is None + batch = _fake_batches(1, 8)[0] + stats = trainer.step(batch, torch.device("cpu"), global_step=1) + assert stats["did_g_step"] is False + + +def test_flow_stage_trainer_ddpm_not_implemented_for_stage2(): + 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"), + )