From 9bf5874308977368c2d2f35e241e714e76e9db60 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 12 Aug 2026 14:35:14 +0200 Subject: [PATCH] Make config dataclasses the single source of truth for DEFAULT_CONFIG DEFAULT_CONFIG and build_models/build_critics/StageSpec.from_config's inline .get(key, default) fallbacks had already drifted: two keys (stage2_model.decoder, stage2_model.particle_type.target) resolved differently depending on whether a config dict came from merge_cli_overrides (fully populated, correct) or was hand-built and partial (fell back to stale v0.2-shaped literals). Introduce frozen dataclasses (GiantConfig and its nested blocks) in giant/config.py as the actual single declaration of every default; DEFAULT_CONFIG is now generated from them instead of hand-maintained, and build_models, build_critics, and StageSpec.from_config consume the dataclasses instead of duplicating literal fallbacks, so this class of drift can't recur. Router/n_sec sub-blocks keep an `extra` catch-all for their genuinely dynamic keys (composed-router axes, runtime-seeded centers_init, legacy_owner). Fixing the fallback surfaced the same latent bug in two existing partial-config callers that had been silently depending on it: a test fixture in test_train.py and scripts/warm_setup_cache.py's minimal cfg (now merged against DEFAULT_CONFIG instead of hand-rolled, closing the gap for good). See issues.md Issue 1. Co-Authored-By: Claude Sonnet 5 --- giant/config.py | 952 ++++++++++++++++++++++++++---------- giant/model/network.py | 120 ++--- giant/pipeline.py | 4 +- giant/training/trainers.py | 77 +-- scripts/warm_setup_cache.py | 27 +- tests/test_config.py | 74 +++ tests/test_network.py | 56 +++ tests/test_train.py | 22 +- 8 files changed, 963 insertions(+), 369 deletions(-) diff --git a/giant/config.py b/giant/config.py index 92c42d9..66d8cd7 100644 --- a/giant/config.py +++ b/giant/config.py @@ -4,6 +4,7 @@ import random import subprocess import sys import tomllib +from dataclasses import dataclass, field from datetime import datetime, timezone from enum import Enum from pathlib import Path @@ -29,272 +30,691 @@ class Conditioning(str, Enum): CONFIG_VERSION = 3 -DEFAULT_CONFIG: dict = { - "conditioning": { - # Width of the fused conditioning vector produced by the encoder's - # fusion MLP, consumed by every downstream trunk. - "out_dim": 128, - # false: stage 1 and stage 2 each construct their own ConditionEncoder - # with identical config but independent weights. true: one instance, - # shared by reference (halves the conditioning parameter count, - # forces a common representation). - "share_stages": False, - "particle": { - # "physical": a small MLP over log(mass)/charge, computable for - # any PDG code — generalizes beyond the training menu. - # "embedding": a learned nn.Embedding over a dense training-vocab - # index — memorizes the training menu; the generalization- - # comparison baseline, and the only mode compatible with - # stage2_model.particle_type.target = "embedding". - # "onehot": a fixed, unlearned vector — top (emb_dim - 1) PDG - # codes by training-set count, plus one "other" bin. NOT a - # reparameterization of "embedding": the vocabulary cap is the - # real difference. - "type": "physical", - # Width of this axis's vector. Under "onehot" this also sets the - # class count. - "emb_dim": 16, - # Depth of the sub-MLP under "physical". Ignored under - # "embedding"/"onehot". - "n_layers": 1, - }, - "material": { - "type": "physical", - "emb_dim": 16, - "n_layers": 1, - }, - }, - "stage1_model": { - # false skips building/training stage 1 entirely. The resulting - # checkpoint holds only stage 2 and cannot be rolled out. - "active": True, - # "flow": conditional flow matching (~10 ODE steps at inference). - # "ddpm": cosine-schedule diffusion baseline. - # "wgan": WGAN-GP, single forward pass at inference. - "generator": "flow", - # Trunk width — also the width of every expert under a routed trunk. - "hidden_dim": 256, - # Number of ResBlocks in the trunk, and in every expert under a - # routed trunk. - "n_res_blocks": 6, - "dropout": 0.0, - # Weight of this stage's loss in the total when both stages are - # active and non-adversarial. A WGAN stage's adversarial loss drives - # its own optimizer, so `lambda` scales only its non-adversarial - # auxiliary terms. - "lambda": 1.0, - "flow": { - # Width of the SinusoidalEmbedding for the flow time variable. - "time_dim": 64, - }, - "ddpm": { - "time_dim": 64, - "n_steps": 1000, - }, - "wgan": { - "noise_dim": 64, - "n_critic": 5, - "gp_weight": 10.0, - # 0.0 means "inherit train.lr" — not None, since the TOML writer - # has no null literal to round-trip. - "critic_lr": 0.0, - # 0 means "inherit stage1_model.hidden_dim/n_res_blocks" — same - # round-trip-friendly sentinel as critic_lr above. - "critic_hidden_dim": 0, - "critic_n_res_blocks": 0, - }, - "router": { - "enabled": False, - "type": "energy", # selects the Router impl from ROUTER_REGISTRY - "n_experts": 4, - "temperature": 0.5, # energy/pdg-router kwarg - "learn_centers": True, # energy/pdg-router kwarg - # energy-router kwargs: mutually exclusive optional learnable - # gate-sharpness modes. learn_width generalizes the shared - # `temperature` to one learnable width per expert; - # learn_temperature instead makes the single shared `temperature` - # itself learnable. Both are bounded to [width_min_ratio, - # width_max_ratio] * temperature so gate sharpness can't run away - # to a collapse-inducing extreme during training. - "learn_width": False, - "learn_temperature": False, - "width_min_ratio": 0.1, - "width_max_ratio": 10.0, - # Importance-CV^2 load-balancing aux loss weight (Shazeer et al. - # 2017). - "lambda_balance": 0.0, - # Entropy-regularization weight penalizing uniform/collapsed - # gating — a secondary guard against all experts' widths - # co-inflating together, which lambda_balance alone can't see. - "lambda_entropy": 0.0, - # Opt-in straight-through Gumbel-softmax train-time combine - # weights: the training forward pass samples a hard one-hot - # combination (matching eval-time top-1 dispatch exactly) while - # the backward pass still flows smooth gradient to every expert. - "gumbel": False, - "gumbel_tau_start": 1.0, - "gumbel_tau_end": 0.1, - "emb_dim": 8, # process/pdg-router kwarg: own pdg(/mat) embedding width - "hidden_dim": 64, # process-router kwarg: its classifier's hidden width - "lambda_proc": 0.0, # process-router kwarg: supervised process-CE weight - # type = "composed" routes on multiple axes at once (e.g. energy x - # pdg), each with its own expert count/hyperparameters. Axes are - # NOT in these defaults — set them as flat axis{i}_{field} keys - # instead of "n_experts", e.g. axis0_type = "energy", - # axis0_n_experts = 4, axis1_type = "pdg", axis1_n_experts = 3, - # axis1_emb_dim = 8. See giant.model.network._parse_composed_axes. - }, - }, - "stage2_model": { - # false trains stage 1 alone. giant rollout must then refuse the - # checkpoint; giant predict still works. - "active": True, - # "one_shot": predict all k_max slots simultaneously with padded - # slots masked from the loss (v0.2 behaviour). - # "autoregressive": emit one secondary at a time in descending-energy - # order. - "decoder": "autoregressive", - # As stage1_model.generator, but under "autoregressive" this is the - # objective for each token. - "generator": "wgan", - "hidden_dim": 256, - "n_res_blocks": 6, - "dropout": 0.0, - "lambda": 1.0, - # Maximum secondary slots. Under "one_shot" this is the fixed output - # width; under "autoregressive" it is a safety cap on the generation - # loop. - "k_max": 15, - # Width of the projected stage-1 outcome fed into stage 2's - # conditioning. - "context_dim": 64, - # "truth": the ground-truth stage-1 target vector, detached — stage- - # level teacher forcing (v0.2 behaviour). "sampled": stage 1's own - # sampled output, closing the train/inference gap at the cost of a - # sampling pass per batch and a moving target early in training. - "stage1_context": "truth", - "n_sec": { - # "head": a classifier over {0..k_max} on the condition encoding - # alone (no diffusion noise), callable independently at - # inference. - # "stop_token": an EOS-style implicit stop — accepted by the - # schema but not implemented in v0.3.0 (see validate_config). - # "truth": take n_sec from ground truth — standalone stage-2 - # evaluation only, never for rollout. - "mode": "head", - "lambda": 0.1, # cross-entropy weight for the head - }, - "particle_type": { - # The three targets mirror the three conditioning.particle - # modes. "onehot": class logits over conditioning.particle.emb_dim - # classes. "physical": regressed (log mass, charge). "embedding": - # regressed against conditioning's own particle embedding table - # (requires conditioning.particle.type = "embedding"). - "target": "onehot", - "lambda": 1.0, - # How a predicted "other" class becomes a concrete PDG code at - # rollout. "sample": draw from the empirical within-bucket - # distribution recorded at map-build time. "modal": always the - # most common member. "drop": discard the secondary. Read only - # under target = "onehot". - "other_policy": "sample", - }, - "autoregressive": { - # Canonical generation order. Single-valued for now; the key - # exists so an alternative ordering is not a config break. - "order": "energy_desc", - # How token i+1 sees tokens <= i. "markov": previous token plus - # running scalars (remaining energy budget, slot index) — a - # fixed-width summary. "attention": causal self-attention over - # all emitted tokens. - "history": "markov", - # "always": condition on the ground-truth previous secondary - # throughout training. "scheduled": scheduled sampling — - # interpolate toward the model's own prediction. "never": - # free-running from the start. - "teacher_forcing": "always", - "tf_p_start": 1.0, - "tf_p_end": 1.0, - "attn_n_heads": 4, - "attn_n_layers": 2, - }, - "flow": { - "time_dim": 64, - }, - "ddpm": { - "time_dim": 64, - "n_steps": 1000, - }, - "wgan": { - "noise_dim": 64, - "n_critic": 5, - "gp_weight": 10.0, - "critic_lr": 0.0, - "critic_hidden_dim": 0, - "critic_n_res_blocks": 0, - # Straight-through Gumbel temperature for the particle-type - # one-hot (distinct from router.gumbel_tau_start/_end, which - # anneal expert-combination weights). Read only under - # particle_type.target = "onehot". - "gumbel_tau_start": 1.0, - "gumbel_tau_end": 0.1, - }, - "router": { - # true: stage 2 shares stage 1's Router module instance, so - # expert i in stage 1 and expert i in stage 2 gate on identical - # conditions by construction — every other key in this block is - # then ignored. false: an independent router. - "tie_to_stage1": False, - "enabled": False, - "type": "energy", - "n_experts": 4, - "temperature": 0.5, - "learn_centers": True, - "learn_width": False, - "learn_temperature": False, - "width_min_ratio": 0.1, - "width_max_ratio": 10.0, - "lambda_balance": 0.0, - "lambda_entropy": 0.0, - "lambda_proc": 0.0, - "gumbel": False, - "gumbel_tau_start": 1.0, - "gumbel_tau_end": 0.1, - "emb_dim": 8, - "hidden_dim": 64, - }, - }, - "train": { - "epochs": 100, - "batch_size": 4096, - "lr": 3e-4, - "weight_decay": 0.01, # AdamW default — exposed so it can be tuned - "ema_decay": 0.9999, # EMA of model weights for sampling; 0 disables - "warmup_epochs": 5, - "val_fraction": 0.1, - # per-epoch val loss (not the marginal/KL validate_every pass) is - # capped to this many batches; 0 = full val set every epoch - "max_val_batches": 200, - "num_workers": 4, - "seed": 0, - "validate_every": 10, - "validate_steps": 10, - # Weights & Biases per-epoch metric logging. Default true in v0.3.0 - # (was opt-in false): the v0.3.0 work is a sequence of architecture - # comparisons, and a run that wasn't logged isn't comparable. Set - # false for throwaway/debug runs. - "wandb": True, - "wandb_project": "giant", - # "" means "use the checkpoint out_dir name" — not None, since the - # TOML writer has no null literal to round-trip. - "wandb_run_name": "", - # Batch-granularity metrics (loss/grad_norm/lr) are logged every N - # optimizer steps, not every batch — a single epoch can be tens of - # thousands of steps. Per-epoch metrics (the metrics.csv row) always - # log in full. - "wandb_log_every": 50, - }, -} +# --- Config dataclasses ----------------------------------------------------- +# +# These are the single source of truth for every default below. DEFAULT_CONFIG +# (a plain dict, for merge_cli_overrides/save_config/TOML round-tripping) is +# *generated* from GiantConfig().to_dict() rather than hand-maintained, so it +# cannot drift from the fallback defaults that build_models/build_critics +# (giant/model/network.py) and StageSpec.from_config (giant/training/trainers.py) +# read off these same dataclasses — see issues.md Issue 1. +# +# Each dataclass is frozen and carries an explicit from_dict/to_dict pair +# (mirroring StageSpec's established style in trainers.py) rather than a +# generic reflection-based helper, so every default is fully type-checkable. +# `lambda` is a Python keyword, so dict key "lambda" is always exposed as the +# field `lambda_weight`. +# +# Two sub-blocks — router and n_sec — carry genuinely dynamic keys that don't +# fit a fixed schema: composed-router `axis{i}_{field}` flags (see +# giant.model.network._parse_composed_axes) and pipeline.py's runtime-seeded +# `centers_init`, plus n_sec's `legacy_owner` (injected only by +# _migrate_legacy_model_config for v0.2 checkpoints). Both dataclasses carry +# an `extra: dict` catch-all so these keys round-trip losslessly without +# becoming named fields that would leak into every new run's config.toml. + + +@dataclass(frozen=True) +class ConditioningAxisConfig: + """One conditioning axis: `conditioning.particle` or `conditioning.material`.""" + + # "physical": a small MLP over log(mass)/charge, computable for any PDG + # code — generalizes beyond the training menu. + # "embedding": a learned nn.Embedding over a dense training-vocab index — + # memorizes the training menu; the generalization-comparison baseline, + # and the only mode compatible with stage2_model.particle_type.target = + # "embedding". + # "onehot": a fixed, unlearned vector — top (emb_dim - 1) PDG codes by + # training-set count, plus one "other" bin. NOT a reparameterization of + # "embedding": the vocabulary cap is the real difference. + type: str = "physical" + # Width of this axis's vector. Under "onehot" this also sets the class + # count. + emb_dim: int = 16 + # Depth of the sub-MLP under "physical". Ignored under "embedding"/"onehot". + n_layers: int = 1 + + @classmethod + def from_dict(cls, d: dict | None) -> "ConditioningAxisConfig": + d = d or {} + return cls( + type=d.get("type", "physical"), + emb_dim=d.get("emb_dim", 16), + n_layers=d.get("n_layers", 1), + ) + + def to_dict(self) -> dict: + return {"type": self.type, "emb_dim": self.emb_dim, "n_layers": self.n_layers} + + +@dataclass(frozen=True) +class ConditioningConfig: + # Width of the fused conditioning vector produced by the encoder's fusion + # MLP, consumed by every downstream trunk. + out_dim: int = 128 + # false: stage 1 and stage 2 each construct their own ConditionEncoder + # with identical config but independent weights. true: one instance, + # shared by reference (halves the conditioning parameter count, forces a + # common representation). + share_stages: bool = False + particle: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig) + material: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "ConditioningConfig": + d = d or {} + return cls( + out_dim=d.get("out_dim", 128), + share_stages=d.get("share_stages", False), + particle=ConditioningAxisConfig.from_dict(d.get("particle")), + material=ConditioningAxisConfig.from_dict(d.get("material")), + ) + + def to_dict(self) -> dict: + return { + "out_dim": self.out_dim, + "share_stages": self.share_stages, + "particle": self.particle.to_dict(), + "material": self.material.to_dict(), + } + + +@dataclass(frozen=True) +class FlowConfig: + # Width of the SinusoidalEmbedding for the flow time variable. + time_dim: int = 64 + + @classmethod + def from_dict(cls, d: dict | None) -> "FlowConfig": + d = d or {} + return cls(time_dim=d.get("time_dim", 64)) + + def to_dict(self) -> dict: + return {"time_dim": self.time_dim} + + +@dataclass(frozen=True) +class DdpmConfig: + time_dim: int = 64 + n_steps: int = 1000 + + @classmethod + def from_dict(cls, d: dict | None) -> "DdpmConfig": + d = d or {} + return cls(time_dim=d.get("time_dim", 64), n_steps=d.get("n_steps", 1000)) + + def to_dict(self) -> dict: + return {"time_dim": self.time_dim, "n_steps": self.n_steps} + + +@dataclass(frozen=True) +class Stage1WganConfig: + noise_dim: int = 64 + n_critic: int = 5 + gp_weight: float = 10.0 + # 0.0 means "inherit train.lr" — not None, since the TOML writer has no + # null literal to round-trip. + critic_lr: float = 0.0 + # 0 means "inherit stage1_model.hidden_dim/n_res_blocks" — same + # round-trip-friendly sentinel as critic_lr above. + critic_hidden_dim: int = 0 + critic_n_res_blocks: int = 0 + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage1WganConfig": + d = d or {} + return cls( + noise_dim=d.get("noise_dim", 64), + n_critic=d.get("n_critic", 5), + gp_weight=d.get("gp_weight", 10.0), + critic_lr=d.get("critic_lr", 0.0), + critic_hidden_dim=d.get("critic_hidden_dim", 0), + critic_n_res_blocks=d.get("critic_n_res_blocks", 0), + ) + + def to_dict(self) -> dict: + return { + "noise_dim": self.noise_dim, + "n_critic": self.n_critic, + "gp_weight": self.gp_weight, + "critic_lr": self.critic_lr, + "critic_hidden_dim": self.critic_hidden_dim, + "critic_n_res_blocks": self.critic_n_res_blocks, + } + + +@dataclass(frozen=True) +class Stage2WganConfig(Stage1WganConfig): + # Straight-through Gumbel temperature for the particle-type one-hot + # (distinct from router.gumbel_tau_start/_end, which anneal + # expert-combination weights). Read only under particle_type.target = + # "onehot". + gumbel_tau_start: float = 1.0 + gumbel_tau_end: float = 0.1 + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage2WganConfig": + d = d or {} + return cls( + noise_dim=d.get("noise_dim", 64), + n_critic=d.get("n_critic", 5), + gp_weight=d.get("gp_weight", 10.0), + critic_lr=d.get("critic_lr", 0.0), + critic_hidden_dim=d.get("critic_hidden_dim", 0), + critic_n_res_blocks=d.get("critic_n_res_blocks", 0), + gumbel_tau_start=d.get("gumbel_tau_start", 1.0), + gumbel_tau_end=d.get("gumbel_tau_end", 0.1), + ) + + def to_dict(self) -> dict: + return { + **super().to_dict(), + "gumbel_tau_start": self.gumbel_tau_start, + "gumbel_tau_end": self.gumbel_tau_end, + } + + +# Fixed router fields shared by stage1_model.router and stage2_model.router. +# Composed-router axis{i}_{field} keys and pipeline.py's runtime-seeded +# centers_init are NOT in this set — they land in RouterConfig.extra instead +# (see giant.model.network._parse_composed_axes). +_ROUTER_KNOWN_KEYS = frozenset( + { + "enabled", + "type", + "n_experts", + "temperature", + "learn_centers", + "learn_width", + "learn_temperature", + "width_min_ratio", + "width_max_ratio", + "lambda_balance", + "lambda_entropy", + "gumbel", + "gumbel_tau_start", + "gumbel_tau_end", + "emb_dim", + "hidden_dim", + "lambda_proc", + } +) + + +@dataclass(frozen=True) +class RouterConfig: + """`stage1_model.router`'s fixed fields. `extra` holds any key not named + below — composed-router `axis{i}_{field}` flags and pipeline.py's + runtime-seeded `centers_init` — so from_dict()/to_dict() round-trip + losslessly without this dataclass needing to know about them.""" + + enabled: bool = False + type: str = "energy" # selects the Router impl from ROUTER_REGISTRY + n_experts: int = 4 + temperature: float = 0.5 # energy/pdg-router kwarg + learn_centers: bool = True # energy/pdg-router kwarg + # energy-router kwargs: mutually exclusive optional learnable + # gate-sharpness modes. learn_width generalizes the shared `temperature` + # to one learnable width per expert; learn_temperature instead makes the + # single shared `temperature` itself learnable. Both are bounded to + # [width_min_ratio, width_max_ratio] * temperature so gate sharpness + # can't run away to a collapse-inducing extreme during training. + learn_width: bool = False + learn_temperature: bool = False + width_min_ratio: float = 0.1 + width_max_ratio: float = 10.0 + # Importance-CV^2 load-balancing aux loss weight (Shazeer et al. 2017). + lambda_balance: float = 0.0 + # Entropy-regularization weight penalizing uniform/collapsed gating — a + # secondary guard against all experts' widths co-inflating together, + # which lambda_balance alone can't see. + lambda_entropy: float = 0.0 + # Opt-in straight-through Gumbel-softmax train-time combine weights: the + # training forward pass samples a hard one-hot combination (matching + # eval-time top-1 dispatch exactly) while the backward pass still flows + # smooth gradient to every expert. + gumbel: bool = False + gumbel_tau_start: float = 1.0 + gumbel_tau_end: float = 0.1 + emb_dim: int = 8 # process/pdg-router kwarg: own pdg(/mat) embedding width + hidden_dim: int = 64 # process-router kwarg: its classifier's hidden width + lambda_proc: float = 0.0 # process-router kwarg: supervised process-CE weight + # type = "composed" routes on multiple axes at once (e.g. energy x pdg), + # each with its own expert count/hyperparameters. Axes are NOT in these + # defaults — set them as flat axis{i}_{field} keys instead of + # "n_experts", e.g. axis0_type = "energy", axis0_n_experts = 4, + # axis1_type = "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See + # giant.model.network._parse_composed_axes. + extra: dict = field(default_factory=dict) + + @classmethod + def from_dict(cls, d: dict | None) -> "RouterConfig": + d = d or {} + return cls( + enabled=d.get("enabled", False), + type=d.get("type", "energy"), + n_experts=d.get("n_experts", 4), + temperature=d.get("temperature", 0.5), + learn_centers=d.get("learn_centers", True), + learn_width=d.get("learn_width", False), + learn_temperature=d.get("learn_temperature", False), + width_min_ratio=d.get("width_min_ratio", 0.1), + width_max_ratio=d.get("width_max_ratio", 10.0), + lambda_balance=d.get("lambda_balance", 0.0), + lambda_entropy=d.get("lambda_entropy", 0.0), + gumbel=d.get("gumbel", False), + gumbel_tau_start=d.get("gumbel_tau_start", 1.0), + gumbel_tau_end=d.get("gumbel_tau_end", 0.1), + emb_dim=d.get("emb_dim", 8), + hidden_dim=d.get("hidden_dim", 64), + lambda_proc=d.get("lambda_proc", 0.0), + extra={k: v for k, v in d.items() if k not in _ROUTER_KNOWN_KEYS}, + ) + + def to_dict(self) -> dict: + return { + "enabled": self.enabled, + "type": self.type, + "n_experts": self.n_experts, + "temperature": self.temperature, + "learn_centers": self.learn_centers, + "learn_width": self.learn_width, + "learn_temperature": self.learn_temperature, + "width_min_ratio": self.width_min_ratio, + "width_max_ratio": self.width_max_ratio, + "lambda_balance": self.lambda_balance, + "lambda_entropy": self.lambda_entropy, + "gumbel": self.gumbel, + "gumbel_tau_start": self.gumbel_tau_start, + "gumbel_tau_end": self.gumbel_tau_end, + "emb_dim": self.emb_dim, + "hidden_dim": self.hidden_dim, + "lambda_proc": self.lambda_proc, + **self.extra, + } + + +@dataclass(frozen=True) +class Stage2RouterConfig(RouterConfig): + # true: stage 2 shares stage 1's Router module instance, so expert i in + # stage 1 and expert i in stage 2 gate on identical conditions by + # construction — every other key in this block is then ignored. false: + # an independent router. + tie_to_stage1: bool = False + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage2RouterConfig": + d = d or {} + known = _ROUTER_KNOWN_KEYS | {"tie_to_stage1"} + return cls( + tie_to_stage1=d.get("tie_to_stage1", False), + enabled=d.get("enabled", False), + type=d.get("type", "energy"), + n_experts=d.get("n_experts", 4), + temperature=d.get("temperature", 0.5), + learn_centers=d.get("learn_centers", True), + learn_width=d.get("learn_width", False), + learn_temperature=d.get("learn_temperature", False), + width_min_ratio=d.get("width_min_ratio", 0.1), + width_max_ratio=d.get("width_max_ratio", 10.0), + lambda_balance=d.get("lambda_balance", 0.0), + lambda_entropy=d.get("lambda_entropy", 0.0), + gumbel=d.get("gumbel", False), + gumbel_tau_start=d.get("gumbel_tau_start", 1.0), + gumbel_tau_end=d.get("gumbel_tau_end", 0.1), + emb_dim=d.get("emb_dim", 8), + hidden_dim=d.get("hidden_dim", 64), + lambda_proc=d.get("lambda_proc", 0.0), + extra={k: v for k, v in d.items() if k not in known}, + ) + + def to_dict(self) -> dict: + return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()} + + +@dataclass(frozen=True) +class NSecConfig: + # "head": a classifier over {0..k_max} on the condition encoding alone + # (no diffusion noise), callable independently at inference. + # "stop_token": an EOS-style implicit stop — accepted by the schema but + # not implemented in v0.3.0 (see validate_config). + # "truth": take n_sec from ground truth — standalone stage-2 evaluation + # only, never for rollout. + mode: str = "head" + lambda_weight: float = 0.1 # dict key "lambda" — cross-entropy weight for the head + # Holds "legacy_owner" when injected by _migrate_legacy_model_config + # (v0.2 checkpoints only) — not a user-facing config.toml key. + extra: dict = field(default_factory=dict) + + @property + def legacy_owner(self) -> str | None: + return self.extra.get("legacy_owner") + + @classmethod + def from_dict(cls, d: dict | None) -> "NSecConfig": + d = d or {} + known = {"mode", "lambda"} + return cls( + mode=d.get("mode", "head"), + lambda_weight=d.get("lambda", 0.1), + extra={k: v for k, v in d.items() if k not in known}, + ) + + def to_dict(self) -> dict: + return {"mode": self.mode, "lambda": self.lambda_weight, **self.extra} + + +@dataclass(frozen=True) +class ParticleTypeConfig: + # The three targets mirror the three conditioning.particle modes. + # "onehot": class logits over conditioning.particle.emb_dim classes. + # "physical": regressed (log mass, charge). "embedding": regressed + # against conditioning's own particle embedding table (requires + # conditioning.particle.type = "embedding"). + target: str = "onehot" + lambda_weight: float = 1.0 # dict key "lambda" + # How a predicted "other" class becomes a concrete PDG code at rollout. + # "sample": draw from the empirical within-bucket distribution recorded + # at map-build time. "modal": always the most common member. "drop": + # discard the secondary. Read only under target = "onehot". + other_policy: str = "sample" + + @classmethod + def from_dict(cls, d: dict | None) -> "ParticleTypeConfig": + d = d or {} + return cls( + target=d.get("target", "onehot"), + lambda_weight=d.get("lambda", 1.0), + other_policy=d.get("other_policy", "sample"), + ) + + def to_dict(self) -> dict: + return {"target": self.target, "lambda": self.lambda_weight, "other_policy": self.other_policy} + + +@dataclass(frozen=True) +class AutoregressiveConfig: + # Canonical generation order. Single-valued for now; the key exists so + # an alternative ordering is not a config break. + order: str = "energy_desc" + # How token i+1 sees tokens <= i. "markov": previous token plus running + # scalars (remaining energy budget, slot index) — a fixed-width summary. + # "attention": causal self-attention over all emitted tokens. + history: str = "markov" + # "always": condition on the ground-truth previous secondary throughout + # training. "scheduled": scheduled sampling — interpolate toward the + # model's own prediction. "never": free-running from the start. + teacher_forcing: str = "always" + tf_p_start: float = 1.0 + tf_p_end: float = 1.0 + attn_n_heads: int = 4 + attn_n_layers: int = 2 + + @classmethod + def from_dict(cls, d: dict | None) -> "AutoregressiveConfig": + d = d or {} + return cls( + order=d.get("order", "energy_desc"), + history=d.get("history", "markov"), + teacher_forcing=d.get("teacher_forcing", "always"), + tf_p_start=d.get("tf_p_start", 1.0), + tf_p_end=d.get("tf_p_end", 1.0), + attn_n_heads=d.get("attn_n_heads", 4), + attn_n_layers=d.get("attn_n_layers", 2), + ) + + def to_dict(self) -> dict: + return { + "order": self.order, + "history": self.history, + "teacher_forcing": self.teacher_forcing, + "tf_p_start": self.tf_p_start, + "tf_p_end": self.tf_p_end, + "attn_n_heads": self.attn_n_heads, + "attn_n_layers": self.attn_n_layers, + } + + +@dataclass(frozen=True) +class Stage1ModelConfig: + # false skips building/training stage 1 entirely. The resulting + # checkpoint holds only stage 2 and cannot be rolled out. + active: bool = True + # "flow": conditional flow matching (~10 ODE steps at inference). + # "ddpm": cosine-schedule diffusion baseline. + # "wgan": WGAN-GP, single forward pass at inference. + generator: str = "flow" + # Trunk width — also the width of every expert under a routed trunk. + hidden_dim: int = 256 + # Number of ResBlocks in the trunk, and in every expert under a routed + # trunk. + n_res_blocks: int = 6 + dropout: float = 0.0 + # Weight of this stage's loss in the total when both stages are active + # and non-adversarial. A WGAN stage's adversarial loss drives its own + # optimizer, so `lambda` scales only its non-adversarial auxiliary terms. + lambda_weight: float = 1.0 # dict key "lambda" + flow: FlowConfig = field(default_factory=FlowConfig) + ddpm: DdpmConfig = field(default_factory=DdpmConfig) + wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig) + router: RouterConfig = field(default_factory=RouterConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage1ModelConfig": + d = d or {} + return cls( + active=d.get("active", True), + generator=d.get("generator", "flow"), + hidden_dim=d.get("hidden_dim", 256), + n_res_blocks=d.get("n_res_blocks", 6), + dropout=d.get("dropout", 0.0), + lambda_weight=d.get("lambda", 1.0), + flow=FlowConfig.from_dict(d.get("flow")), + ddpm=DdpmConfig.from_dict(d.get("ddpm")), + wgan=Stage1WganConfig.from_dict(d.get("wgan")), + router=RouterConfig.from_dict(d.get("router")), + ) + + def to_dict(self) -> dict: + return { + "active": self.active, + "generator": self.generator, + "hidden_dim": self.hidden_dim, + "n_res_blocks": self.n_res_blocks, + "dropout": self.dropout, + "lambda": self.lambda_weight, + "flow": self.flow.to_dict(), + "ddpm": self.ddpm.to_dict(), + "wgan": self.wgan.to_dict(), + "router": self.router.to_dict(), + } + + +@dataclass(frozen=True) +class Stage2ModelConfig: + # false trains stage 1 alone. giant rollout must then refuse the + # checkpoint; giant predict still works. + active: bool = True + # "one_shot": predict all k_max slots simultaneously with padded slots + # masked from the loss (v0.2 behaviour). + # "autoregressive": emit one secondary at a time in descending-energy + # order. + decoder: str = "autoregressive" + # As stage1_model.generator, but under "autoregressive" this is the + # objective for each token. + generator: str = "wgan" + hidden_dim: int = 256 + n_res_blocks: int = 6 + dropout: float = 0.0 + lambda_weight: float = 1.0 # dict key "lambda" + # Maximum secondary slots. Under "one_shot" this is the fixed output + # width; under "autoregressive" it is a safety cap on the generation loop. + k_max: int = 15 + # Width of the projected stage-1 outcome fed into stage 2's conditioning. + context_dim: int = 64 + # "truth": the ground-truth stage-1 target vector, detached — stage-level + # teacher forcing (v0.2 behaviour). "sampled": stage 1's own sampled + # output, closing the train/inference gap at the cost of a sampling pass + # per batch and a moving target early in training. + stage1_context: str = "truth" + n_sec: NSecConfig = field(default_factory=NSecConfig) + particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig) + autoregressive: AutoregressiveConfig = field(default_factory=AutoregressiveConfig) + flow: FlowConfig = field(default_factory=FlowConfig) + ddpm: DdpmConfig = field(default_factory=DdpmConfig) + wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig) + router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage2ModelConfig": + d = d or {} + return cls( + active=d.get("active", True), + decoder=d.get("decoder", "autoregressive"), + generator=d.get("generator", "wgan"), + hidden_dim=d.get("hidden_dim", 256), + n_res_blocks=d.get("n_res_blocks", 6), + dropout=d.get("dropout", 0.0), + lambda_weight=d.get("lambda", 1.0), + k_max=d.get("k_max", 15), + context_dim=d.get("context_dim", 64), + stage1_context=d.get("stage1_context", "truth"), + n_sec=NSecConfig.from_dict(d.get("n_sec")), + particle_type=ParticleTypeConfig.from_dict(d.get("particle_type")), + autoregressive=AutoregressiveConfig.from_dict(d.get("autoregressive")), + flow=FlowConfig.from_dict(d.get("flow")), + ddpm=DdpmConfig.from_dict(d.get("ddpm")), + wgan=Stage2WganConfig.from_dict(d.get("wgan")), + router=Stage2RouterConfig.from_dict(d.get("router")), + ) + + def to_dict(self) -> dict: + return { + "active": self.active, + "decoder": self.decoder, + "generator": self.generator, + "hidden_dim": self.hidden_dim, + "n_res_blocks": self.n_res_blocks, + "dropout": self.dropout, + "lambda": self.lambda_weight, + "k_max": self.k_max, + "context_dim": self.context_dim, + "stage1_context": self.stage1_context, + "n_sec": self.n_sec.to_dict(), + "particle_type": self.particle_type.to_dict(), + "autoregressive": self.autoregressive.to_dict(), + "flow": self.flow.to_dict(), + "ddpm": self.ddpm.to_dict(), + "wgan": self.wgan.to_dict(), + "router": self.router.to_dict(), + } + + +@dataclass(frozen=True) +class TrainConfig: + epochs: int = 100 + batch_size: int = 4096 + lr: float = 3e-4 + weight_decay: float = 0.01 # AdamW default — exposed so it can be tuned + ema_decay: float = 0.9999 # EMA of model weights for sampling; 0 disables + warmup_epochs: int = 5 + val_fraction: float = 0.1 + # per-epoch val loss (not the marginal/KL validate_every pass) is capped + # to this many batches; 0 = full val set every epoch + max_val_batches: int = 200 + num_workers: int = 4 + seed: int = 0 + validate_every: int = 10 + validate_steps: int = 10 + # Weights & Biases per-epoch metric logging. Default true in v0.3.0 (was + # opt-in false): the v0.3.0 work is a sequence of architecture + # comparisons, and a run that wasn't logged isn't comparable. Set false + # for throwaway/debug runs. + wandb: bool = True + wandb_project: str = "giant" + # "" means "use the checkpoint out_dir name" — not None, since the TOML + # writer has no null literal to round-trip. + wandb_run_name: str = "" + # Batch-granularity metrics (loss/grad_norm/lr) are logged every N + # optimizer steps, not every batch — a single epoch can be tens of + # thousands of steps. Per-epoch metrics (the metrics.csv row) always log + # in full. + wandb_log_every: int = 50 + + @classmethod + def from_dict(cls, d: dict | None) -> "TrainConfig": + d = d or {} + return cls( + epochs=d.get("epochs", 100), + batch_size=d.get("batch_size", 4096), + lr=d.get("lr", 3e-4), + weight_decay=d.get("weight_decay", 0.01), + ema_decay=d.get("ema_decay", 0.9999), + warmup_epochs=d.get("warmup_epochs", 5), + val_fraction=d.get("val_fraction", 0.1), + max_val_batches=d.get("max_val_batches", 200), + num_workers=d.get("num_workers", 4), + seed=d.get("seed", 0), + validate_every=d.get("validate_every", 10), + validate_steps=d.get("validate_steps", 10), + wandb=d.get("wandb", True), + wandb_project=d.get("wandb_project", "giant"), + wandb_run_name=d.get("wandb_run_name", ""), + wandb_log_every=d.get("wandb_log_every", 50), + ) + + def to_dict(self) -> dict: + return { + "epochs": self.epochs, + "batch_size": self.batch_size, + "lr": self.lr, + "weight_decay": self.weight_decay, + "ema_decay": self.ema_decay, + "warmup_epochs": self.warmup_epochs, + "val_fraction": self.val_fraction, + "max_val_batches": self.max_val_batches, + "num_workers": self.num_workers, + "seed": self.seed, + "validate_every": self.validate_every, + "validate_steps": self.validate_steps, + "wandb": self.wandb, + "wandb_project": self.wandb_project, + "wandb_run_name": self.wandb_run_name, + "wandb_log_every": self.wandb_log_every, + } + + +@dataclass(frozen=True) +class GiantConfig: + """Root config dataclass — the single source of truth for every default + in DEFAULT_CONFIG below, which is generated from `GiantConfig().to_dict()` + rather than hand-maintained (see issues.md Issue 1).""" + + conditioning: ConditioningConfig = field(default_factory=ConditioningConfig) + stage1_model: Stage1ModelConfig = field(default_factory=Stage1ModelConfig) + stage2_model: Stage2ModelConfig = field(default_factory=Stage2ModelConfig) + train: TrainConfig = field(default_factory=TrainConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "GiantConfig": + d = d or {} + return cls( + conditioning=ConditioningConfig.from_dict(d.get("conditioning")), + stage1_model=Stage1ModelConfig.from_dict(d.get("stage1_model")), + stage2_model=Stage2ModelConfig.from_dict(d.get("stage2_model")), + train=TrainConfig.from_dict(d.get("train")), + ) + + def to_dict(self) -> dict: + return { + "conditioning": self.conditioning.to_dict(), + "stage1_model": self.stage1_model.to_dict(), + "stage2_model": self.stage2_model.to_dict(), + "train": self.train.to_dict(), + } + + +DEFAULT_CONFIG: dict = GiantConfig().to_dict() def git_hash() -> str: diff --git a/giant/model/network.py b/giant/model/network.py index bed790c..e946e20 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -7,6 +7,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig from giant.constants import ( COND_DIM, COND_DIM_BASE, @@ -1593,78 +1594,84 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: particle_cfg = conditioning["particle"] material_cfg = conditioning["material"] particle_conditioning = particle_cfg["type"] - s1cfg = cfg["stage1_model"] - s2cfg = cfg["stage2_model"] - cond_out_dim = conditioning.get("out_dim", 128) + conditioning_cfg = ConditioningConfig.from_dict(conditioning) + s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"]) + s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"]) + cond_out_dim = conditioning_cfg.out_dim shared_cond_enc: ConditionEncoder | None = None - if conditioning.get("share_stages"): + if conditioning_cfg.share_stages: shared_cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None} stage1_router: Router | None = None - if s1cfg.get("active", True): - router_cfg = s1cfg.get("router") or {} - if router_cfg.get("enabled"): + if s1_spec.active: + router_cfg = cfg["stage1_model"].get("router") or {} + if s1_spec.router.enabled: stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning) - generator = s1cfg.get("generator", "flow") - gen_sub = s1cfg.get(generator, {}) or {} - legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner") - n_sec_head_k_max = s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None + generator = s1_spec.generator + # wgan has no time_dim concept (no diffusion/flow time variable) — + # matches the pre-dataclass .get("time_dim", 64) fallback, which + # always hit its default for a wgan sub-block too. + time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64 + legacy_owner = s2_spec.n_sec.legacy_owner + n_sec_head_k_max = s2_spec.k_max if legacy_owner == "stage1" else None result["stage1"] = Stage1Model( pdg_vocab=pdg_vocab, mat_vocab=mat_vocab, particle_cfg=particle_cfg, material_cfg=material_cfg, - hidden_dim=s1cfg.get("hidden_dim", 256), - n_res_blocks=s1cfg.get("n_res_blocks", 6), + hidden_dim=s1_spec.hidden_dim, + n_res_blocks=s1_spec.n_res_blocks, cond_out_dim=cond_out_dim, - dropout=s1cfg.get("dropout", 0.0), + dropout=s1_spec.dropout, generator=generator, - time_dim=gen_sub.get("time_dim", 64), - noise_dim=(s1cfg.get("wgan") or {}).get("noise_dim", 64), + time_dim=time_dim, + noise_dim=s1_spec.wgan.noise_dim, router=stage1_router, n_sec_head_k_max=n_sec_head_k_max, cond_enc=shared_cond_enc, ) - if s2cfg.get("active", True): - decoder = s2cfg.get("decoder", "one_shot") - router_cfg = s2cfg.get("router") or {} + if s2_spec.active: + decoder = s2_spec.decoder + router_cfg = cfg["stage2_model"].get("router") or {} stage2_router: Router | None = None - if router_cfg.get("enabled"): - if router_cfg.get("tie_to_stage1") and stage1_router is not None: + if s2_spec.router.enabled: + if s2_spec.router.tie_to_stage1 and stage1_router is not None: stage2_router = stage1_router else: stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning) - generator = s2cfg.get("generator", "wgan") - gen_sub = s2cfg.get(generator, {}) or {} - legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner") - k_max = s2cfg.get("k_max", K_MAX) - particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"} + generator = s2_spec.generator + # wgan has no time_dim concept — see the matching comment in stage 1 + # above. + time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64 + legacy_owner = s2_spec.n_sec.legacy_owner + k_max = s2_spec.k_max + particle_type_cfg = s2_spec.particle_type.to_dict() if decoder == "autoregressive": - ar_cfg = s2cfg.get("autoregressive") or {} + ar_cfg = s2_spec.autoregressive result["stage2"] = Stage2Autoregressive( pdg_vocab=pdg_vocab, mat_vocab=mat_vocab, particle_cfg=particle_cfg, material_cfg=material_cfg, - hidden_dim=s2cfg.get("hidden_dim", 256), - n_res_blocks=s2cfg.get("n_res_blocks", 6), + hidden_dim=s2_spec.hidden_dim, + n_res_blocks=s2_spec.n_res_blocks, cond_out_dim=cond_out_dim, - context_dim=s2cfg.get("context_dim", 64), - dropout=s2cfg.get("dropout", 0.0), + context_dim=s2_spec.context_dim, + dropout=s2_spec.dropout, generator=generator, - time_dim=gen_sub.get("time_dim", 64), - noise_dim=(s2cfg.get("wgan") or {}).get("noise_dim", 64), + time_dim=time_dim, + noise_dim=s2_spec.wgan.noise_dim, k_max=k_max, router=stage2_router, build_n_sec_head=legacy_owner != "stage1", particle_type_cfg=particle_type_cfg, - history=ar_cfg.get("history", "markov"), - attn_n_heads=ar_cfg.get("attn_n_heads", 4), - attn_n_layers=ar_cfg.get("attn_n_layers", 2), + history=ar_cfg.history, + attn_n_heads=ar_cfg.attn_n_heads, + attn_n_layers=ar_cfg.attn_n_layers, cond_enc=shared_cond_enc, ) else: @@ -1674,15 +1681,15 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: mat_vocab=mat_vocab, particle_cfg=particle_cfg, material_cfg=material_cfg, - hidden_dim=s2cfg.get("hidden_dim", 256), - n_res_blocks=s2cfg.get("n_res_blocks", 6), + hidden_dim=s2_spec.hidden_dim, + n_res_blocks=s2_spec.n_res_blocks, cond_out_dim=cond_out_dim, - context_dim=s2cfg.get("context_dim", 64), + context_dim=s2_spec.context_dim, sec_dim=sec_dim, - dropout=s2cfg.get("dropout", 0.0), + dropout=s2_spec.dropout, generator=generator, - time_dim=gen_sub.get("time_dim", 64), - noise_dim=(s2cfg.get("wgan") or {}).get("noise_dim", 64), + time_dim=time_dim, + noise_dim=s2_spec.wgan.noise_dim, k_max=k_max, router=stage2_router, build_n_sec_head=legacy_owner != "stage1", @@ -1704,29 +1711,30 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: conditioning = cfg["conditioning"] particle_cfg = conditioning["particle"] material_cfg = conditioning["material"] - cond_out_dim = conditioning.get("out_dim", 128) - s1cfg = cfg["stage1_model"] - s2cfg = cfg["stage2_model"] + conditioning_cfg = ConditioningConfig.from_dict(conditioning) + cond_out_dim = conditioning_cfg.out_dim + s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"]) + s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"]) result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None} - if s1cfg.get("active", True) and s1cfg.get("generator") == "wgan": + if s1_spec.active and s1_spec.generator == "wgan": result["stage1"] = CriticModel( pdg_vocab=pdg_vocab, mat_vocab=mat_vocab, particle_cfg=particle_cfg, material_cfg=material_cfg, in_dim=X_DIM, - hidden_dim=s1cfg.get("hidden_dim", 256), - n_res_blocks=s1cfg.get("n_res_blocks", 6), + hidden_dim=s1_spec.hidden_dim, + n_res_blocks=s1_spec.n_res_blocks, cond_out_dim=cond_out_dim, - dropout=s1cfg.get("dropout", 0.0), + dropout=s1_spec.dropout, stage="stage1", ) - if s2cfg.get("active", True) and s2cfg.get("generator") == "wgan": - k_max = s2cfg.get("k_max", K_MAX) - particle_type_cfg = s2cfg.get("particle_type") or {"target": "physical"} + if s2_spec.active and s2_spec.generator == "wgan": + k_max = s2_spec.k_max + particle_type_cfg = s2_spec.particle_type.to_dict() in_dim = stage2_trunk_sec_dim(particle_type_cfg, "wgan", k_max, particle_cfg["emb_dim"]) result["stage2"] = CriticModel( pdg_vocab=pdg_vocab, @@ -1734,12 +1742,12 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]: particle_cfg=particle_cfg, material_cfg=material_cfg, in_dim=in_dim, - hidden_dim=s2cfg.get("hidden_dim", 256), - n_res_blocks=s2cfg.get("n_res_blocks", 6), + hidden_dim=s2_spec.hidden_dim, + n_res_blocks=s2_spec.n_res_blocks, cond_out_dim=cond_out_dim, - dropout=s2cfg.get("dropout", 0.0), + dropout=s2_spec.dropout, stage="stage2", - context_dim=s2cfg.get("context_dim", 64), + context_dim=s2_spec.context_dim, ) return result diff --git a/giant/pipeline.py b/giant/pipeline.py index b6367ee..683053a 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -182,7 +182,7 @@ def run_setup_stage( # if both consumers are active. The material axis is independent. particle_cfg = cfg["conditioning"]["particle"] material_cfg = cfg["conditioning"]["material"] - particle_type_target = cfg["stage2_model"].get("particle_type", {}).get("target") + particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target pdg_topn_map: TopNMap | None = None if particle_cfg["type"] == "onehot" or particle_type_target == "onehot": @@ -379,7 +379,7 @@ def run_train_job( # The secondary type-index map depends on stage2_model.particle_type.target, # independently of conditioning's own onehot/embedding choice above # (physical stays untouched/None). - particle_type_target = cfg["stage2_model"].get("particle_type", {}).get("target", "physical") + particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target if particle_type_target == "onehot": assert setup.pdg_topn_map is not None sec_type_class_map = setup.pdg_topn_map.class_map diff --git a/giant/training/trainers.py b/giant/training/trainers.py index 8473a87..d6ffd6d 100644 --- a/giant/training/trainers.py +++ b/giant/training/trainers.py @@ -21,6 +21,7 @@ import torch import torch.nn.functional as F import torch.optim as optim +from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig from giant.constants import CONT_SLOT_DIM from giant.model.network import Router, stage2_type_dim from giant.model.schedule import ( @@ -94,7 +95,7 @@ class StageSpec: n_sec_lambda: float = 0.1 # particle-type target (stage 2 only) - particle_type: dict = field(default_factory=lambda: {"target": "physical"}) + particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig) particle_type_emb_dim: int = 16 # optimization @@ -128,48 +129,56 @@ class StageSpec: @classmethod def from_config(cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int) -> "StageSpec": - t = cfg["train"] - stage_cfg = cfg[f"{name}_model"] - router_cfg = stage_cfg.get("router") or {} - wgan_cfg = stage_cfg.get("wgan") or {} - ar_cfg = (stage_cfg.get("autoregressive") or {}) if is_stage2 else {} + t = TrainConfig.from_dict(cfg["train"]) + # n_sec/particle_type/decoder/autoregressive/wgan's gumbel_tau_* are + # stage-2-only concepts, always read off s2_spec (guarded by + # is_stage2 where the stage-1 StageSpec needs a different value) — + # historically n_sec/particle_type were read from stage2_model + # unconditionally even for the stage-1 StageSpec, preserved here for + # behavioral parity. stage_spec covers the fields both stage configs + # share structurally (generator, lambda, router, ddpm, and wgan's + # base fields — Stage2ModelConfig's sub-configs all subclass + # stage 1's). + s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"]) + stage_spec = s2_spec if is_stage2 else Stage1ModelConfig.from_dict(cfg["stage1_model"]) return cls( name=name, is_stage2=is_stage2, - generator=stage_cfg["generator"], - decoder=stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot", - lambda_weight=stage_cfg.get("lambda", 1.0), - n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1), - particle_type=cfg["stage2_model"].get("particle_type") or {"target": "physical"}, + generator=stage_spec.generator, + decoder=s2_spec.decoder if is_stage2 else "one_shot", + lambda_weight=stage_spec.lambda_weight, + n_sec_lambda=s2_spec.n_sec.lambda_weight, + particle_type=s2_spec.particle_type, particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"], # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge - # (giant/config.py), so they read directly; the field defaults - # below exist only for tests that construct StageSpec by hand. - lr=t["lr"], - weight_decay=t["weight_decay"], - ema_decay=t["ema_decay"], - warmup_epochs=t["warmup_epochs"], - epochs=t["epochs"], + # (giant/config.py), so TrainConfig.from_dict never has to fall + # back to a literal here; the field defaults below exist only + # for tests that construct StageSpec by hand. + lr=t.lr, + weight_decay=t.weight_decay, + ema_decay=t.ema_decay, + warmup_epochs=t.warmup_epochs, + epochs=t.epochs, steps_per_epoch=max(steps_per_epoch, 1), - 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), - teacher_forcing=ar_cfg.get("teacher_forcing", "always"), - tf_p_start=ar_cfg.get("tf_p_start", 1.0), - tf_p_end=ar_cfg.get("tf_p_end", 1.0), + lambda_balance=stage_spec.router.lambda_balance, + lambda_proc=stage_spec.router.lambda_proc, + lambda_entropy=stage_spec.router.lambda_entropy, + gumbel_tau_start=stage_spec.router.gumbel_tau_start, + gumbel_tau_end=stage_spec.router.gumbel_tau_end, + teacher_forcing=s2_spec.autoregressive.teacher_forcing if is_stage2 else cls.teacher_forcing, + tf_p_start=s2_spec.autoregressive.tf_p_start if is_stage2 else cls.tf_p_start, + tf_p_end=s2_spec.autoregressive.tf_p_end if is_stage2 else cls.tf_p_end, # AR self-sampling under scheduled/never teacher forcing reuses # train.validate_steps as its flow-matching ODE step count — no # dedicated config key for this (the autoregressive config lists # tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only). - ar_sample_steps=t["validate_steps"], - ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000), - n_critic=wgan_cfg.get("n_critic", 5), - gp_weight=wgan_cfg.get("gp_weight", 10.0), - critic_lr=wgan_cfg.get("critic_lr", 0.0), - type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0), - type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1), + ar_sample_steps=t.validate_steps, + ddpm_n_steps=stage_spec.ddpm.n_steps, + n_critic=stage_spec.wgan.n_critic, + gp_weight=stage_spec.wgan.gp_weight, + critic_lr=stage_spec.wgan.critic_lr, + type_gumbel_tau_start=s2_spec.wgan.gumbel_tau_start if is_stage2 else cls.type_gumbel_tau_start, + type_gumbel_tau_end=s2_spec.wgan.gumbel_tau_end if is_stage2 else cls.type_gumbel_tau_end, ) @@ -227,7 +236,7 @@ class StageTrainer: self.router = _stage_router(self.model) self._modules = (self.model, *extra_modules) - self.particle_type_cfg = dict(spec.particle_type or {"target": "physical"}) + self.particle_type_cfg = spec.particle_type.to_dict() self.particle_type_emb_dim = spec.particle_type_emb_dim self.ema_decay = spec.ema_decay diff --git a/scripts/warm_setup_cache.py b/scripts/warm_setup_cache.py index 3168b9e..1feed1c 100644 --- a/scripts/warm_setup_cache.py +++ b/scripts/warm_setup_cache.py @@ -9,6 +9,7 @@ for the sidecar itself. from pathlib import Path +from giant import config as gconfig from giant.constants import K_MAX from giant.pipeline import run_setup_stage @@ -43,19 +44,25 @@ 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,material}.type, stage{1,2}_model.router). This - # CLI only ever configures one router (matching today's single + # Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so + # run_setup_stage always sees every key it might read (e.g. + # conditioning.particle.emb_dim, stage2_model.particle_type.target) at + # its real default, not silently missing/None — see issues.md Issue 1. + # 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": particle_conditioning}, - "material": {"type": material_conditioning}, + cfg = gconfig.merge_cli_overrides( + gconfig.DEFAULT_CONFIG, + None, + { + "conditioning": { + "particle": {"type": particle_conditioning}, + "material": {"type": material_conditioning}, + }, + "stage1_model": {"router": router_cfg}, + "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX}, }, - "stage1_model": {"router": router_cfg}, - "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX}, - } + ) run_setup_stage( Path(data), val_fraction=val_fraction, diff --git a/tests/test_config.py b/tests/test_config.py index 6aea480..cecb74b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -27,6 +27,80 @@ def test_conditioning_enum_has_onehot(): } +# --------------------------------------------------------------------------- +# Config dataclasses (issues.md Issue 1) +# --------------------------------------------------------------------------- + + +def test_giant_config_to_dict_matches_default_config(): + """DEFAULT_CONFIG is generated from GiantConfig().to_dict() (not + hand-maintained), so the two cannot structurally drift apart — but this + pins the *equality* too, catching e.g. a stray in-place mutation of + DEFAULT_CONFIG added elsewhere after import.""" + assert gconfig.GiantConfig().to_dict() == gconfig.DEFAULT_CONFIG + + +@pytest.mark.parametrize( + "cls", + [ + gconfig.ConditioningAxisConfig, + gconfig.ConditioningConfig, + gconfig.FlowConfig, + gconfig.DdpmConfig, + gconfig.Stage1WganConfig, + gconfig.Stage2WganConfig, + gconfig.RouterConfig, + gconfig.Stage2RouterConfig, + gconfig.NSecConfig, + gconfig.ParticleTypeConfig, + gconfig.AutoregressiveConfig, + gconfig.Stage1ModelConfig, + gconfig.Stage2ModelConfig, + gconfig.TrainConfig, + gconfig.GiantConfig, + ], +) +def test_config_dataclass_from_dict_round_trips_through_to_dict(cls): + assert cls.from_dict(cls().to_dict()) == cls() + assert cls.from_dict(None) == cls() + + +def test_stage2_model_config_defaults_match_documented_v030_intent(): + """The two keys issues.md Issue 1 found drifted between DEFAULT_CONFIG + and build_models/StageSpec.from_config's own .get(key, default) + fallbacks — pinned directly against the dataclass that is now their + shared single source of truth.""" + spec = gconfig.Stage2ModelConfig() + assert spec.decoder == "autoregressive" + assert spec.particle_type.target == "onehot" + + +def test_router_config_extra_round_trips_composed_axis_keys(): + d = {"enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4} + router = gconfig.RouterConfig.from_dict(d) + assert router.enabled is True + assert router.extra == {"axis0_type": "energy", "axis0_n_experts": 4} + assert router.to_dict()["axis0_type"] == "energy" + + +def test_stage2_router_config_tie_to_stage1_not_leaked_into_extra(): + router = gconfig.Stage2RouterConfig.from_dict({"tie_to_stage1": True}) + assert router.tie_to_stage1 is True + assert "tie_to_stage1" not in router.extra + + +def test_stage1_router_config_has_no_tie_to_stage1_key(): + """Stage 1's router schema must not gain stage 2's tie_to_stage1 key — + that would change every future run's saved config.toml shape.""" + assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict() + + +def test_n_sec_config_extra_round_trips_legacy_owner(): + n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "legacy_owner": "stage1"}) + assert n_sec.legacy_owner == "stage1" + assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "legacy_owner": "stage1"} + + # --------------------------------------------------------------------------- # _deep_merge # --------------------------------------------------------------------------- diff --git a/tests/test_network.py b/tests/test_network.py index 5deab01..7857d39 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -12,6 +12,7 @@ from giant.model.network import ( Stage1Model, Stage2Autoregressive, Stage2OneShot, + build_critics, build_models, cat_col_layout, stage2_trunk_sec_dim, @@ -652,3 +653,58 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete assert shared_ids assert shared_ids <= {id(p) for p in stage1.parameters()} assert shared_ids <= {id(p) for p in stage2.parameters()} + + +# ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─ + + +def _partial_model_config() -> dict: + """A hand-built model_config that omits stage2_model.decoder and + stage2_model.particle_type — deliberately not derived from + DEFAULT_CONFIG, unlike _minimal_model_config above. Regression fixture + for issues.md Issue 1: build_models/build_critics/StageSpec.from_config's + own fallback defaults for these two keys must equal DEFAULT_CONFIG's + ("autoregressive" / "onehot"), not the old, now-wrong v0.2-shaped + ("one_shot" / "physical") literals that used to live in three separate + .get(key, default) call sites.""" + return { + "pdg_vocab": 3, + "mat_vocab": 2, + "conditioning": { + "particle": {"type": "physical", "emb_dim": 4, "n_layers": 1}, + "material": {"type": "physical", "emb_dim": 4, "n_layers": 1}, + }, + "stage1_model": {"active": False}, + "stage2_model": { + "generator": "wgan", + "hidden_dim": 8, + "n_res_blocks": 1, + "k_max": 3, + # decoder and particle_type deliberately omitted + }, + } + + +def test_build_models_omitted_decoder_and_particle_type_match_default_config(): + built = build_models(_partial_model_config()) + assert isinstance(built["stage2"], Stage2Autoregressive) + assert built["stage2"].particle_type_cfg["target"] == "onehot" + + +def test_build_critics_omitted_particle_type_matches_default_config(): + cfg = _partial_model_config() + cfg["stage2_model"]["generator"] = "wgan" + onehot_critic = build_critics(cfg)["stage2"] + assert onehot_critic is not None + onehot_in_dim = onehot_critic.input_proj.in_features + + cfg["stage2_model"]["particle_type"] = {"target": "physical"} + physical_critic = build_critics(cfg)["stage2"] + assert physical_critic is not None + physical_in_dim = physical_critic.input_proj.in_features + + # onehot's per-slot type width is emb_dim classes vs. physical's fixed + # (log-mass, charge) pair — different unless emb_dim happens to be 2, so + # this also confirms the critic was actually built in onehot mode by + # default, not silently falling back to physical. + assert onehot_in_dim != physical_in_dim diff --git a/tests/test_train.py b/tests/test_train.py index 9a220fe..fc5781d 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -265,6 +265,12 @@ def _base_cfg(): "k_max": K_MAX, "context_dim": 16, "n_sec": {"mode": "head", "lambda": 0.1}, + # Explicit, not relying on the fallback default (which is + # "onehot", matching DEFAULT_CONFIG — see issues.md Issue 1): + # the "physical"-labelled cases below (and this fixture's own + # comment history) intend this as the base "physical" case, + # with "*_onehot"/"*_embedding" cases opting in explicitly. + "particle_type": {"target": "physical", "lambda": 1.0}, "flow": {"time_dim": 16}, "ddpm": {"time_dim": 16, "n_steps": 50}, "wgan": { @@ -566,6 +572,20 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2(): FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu")) +def test_stage_spec_from_config_omitted_decoder_and_particle_type_match_default_config(): + """Regression for issues.md Issue 1: StageSpec.from_config's own fallback + defaults for stage2_model.decoder/particle_type must equal + DEFAULT_CONFIG's ("autoregressive" / "onehot"), not the old, now-wrong + ("one_shot" / "physical") literals a .get(key, default) call used to + supply when a hand-built cfg omitted these keys.""" + cfg = _base_cfg() + del cfg["stage2_model"]["decoder"] + del cfg["stage2_model"]["particle_type"] + spec = StageSpec.from_config(cfg, "stage2", is_stage2=True, steps_per_epoch=1) + assert spec.decoder == "autoregressive" + assert spec.particle_type.target == "onehot" + + # --- AR trainer wiring (v0.3.0 step 5) -------------------------------------- @@ -664,7 +684,7 @@ def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation(): def test_wgan_physical_omits_grad_norm_slice_columns(): - cfg = _base_cfg() # default stage2_model has no particle_type -> "physical" + cfg = _base_cfg() # _base_cfg's stage2_model.particle_type.target is "physical" with tempfile.TemporaryDirectory() as tmp: out_dir = Path(tmp) / "run" _run_train(cfg, out_dir)