diff --git a/giant/config.py b/giant/config.py index e4945e5..77814af 100644 --- a/giant/config.py +++ b/giant/config.py @@ -1,3 +1,4 @@ +import copy import hashlib import random import subprocess @@ -12,133 +13,287 @@ import torch class Conditioning(str, Enum): - """`model.conditioning` choices — shared by `giant.cli` and `scripts.dwarf`'s - Typer commands so the two CLIs can't silently drift apart on the option's - valid values (see DEFAULT_CONFIG["model"]["conditioning"] for what each - value means).""" + """`conditioning.particle.type` / `conditioning.material.type` choices — + shared by `giant.cli` and `scripts.dwarf`'s Typer commands so the two + CLIs can't silently drift apart on the option's valid values (see + DEFAULT_CONFIG["conditioning"] for what each value means).""" physical = "physical" embedding = "embedding" + onehot = "onehot" + + +# Tags a config dict (config.toml, or a checkpoint's model_config) as the new +# v0.3 nested format. Absence of `[meta].config_version == CONFIG_VERSION` is +# read as "this is a v0.2 dict" by migrate_config below. +CONFIG_VERSION = 3 DEFAULT_CONFIG: dict = { - "train": { - "mode": "flow", - "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 - # 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, - "val_fraction": 0.1, - "num_workers": 4, - "seed": 0, - "validate_every": 10, - "validate_steps": 10, - "warmup_epochs": 5, - "lambda_nsec": 0.1, - "lambda_s2": 1.0, - # WGAN-GP-only knobs (mode == "wgan"; ignored by flow/ddpm). n_critic: - # critic updates per generator update. gp_weight: gradient-penalty - # coefficient (Gulrajani et al. 2017). critic_lr: 0.0 means "use - # `lr`" — not None, since save_config's TOML writer has no null - # literal to round-trip. - "n_critic": 5, - "gp_weight": 10.0, - "critic_lr": 0.0, - # Weights & Biases per-epoch metric logging (opt-in; see giant.train). - # "" for wandb_run_name means "use the checkpoint out_dir name" — not - # None, since save_config's TOML writer has no null literal to - # round-trip. - "wandb": False, - "wandb_project": "giant", - "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 (see steps_per_epoch above), and logging every - # one of them would flood the run with points the UI has to downsample - # anyway. Per-epoch metrics (the metrics.csv row) always log in full. - "wandb_log_every": 50, + "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, + }, }, - "model": { + "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, - "n_blocks": 6, - "emb_dim": 16, - "dropout": 0.1, - # WGAN generator noise-vector width (mode == "wgan" only). - "noise_dim": 64, - # "physical" conditions on material/particle physical properties via - # a small MLP (giant.model.network.ConditionEncoder); "embedding" - # keeps the original learned pdg/material embedding tables — kept - # available as the generalization-comparison baseline. Checkpoints - # from before this option existed have no "conditioning" key and - # load as "embedding" (see giant.model.network.build_models). - "conditioning": "physical", + # 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, - # 0 means "inherit model.hidden_dim/n_blocks" (see - # resolve_expert_dims below) — not a fixed 128/3, which silently - # ignored --hidden-dim/--n-blocks whenever routing was enabled. - # TOML has no null literal to round-trip (same pattern as - # critic_lr/wandb_run_name above), hence 0 rather than None. - "expert_hidden_dim": 0, - "expert_n_blocks": 0, "temperature": 0.5, # energy/pdg-router kwarg "learn_centers": True, # energy/pdg-router kwarg # energy-router kwargs: mutually exclusive optional learnable - # gate-sharpness modes (see giant.model.network.EnergyRouter). - # 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 - # (sigmoid-parameterized, warm-started to reproduce `temperature` - # exactly at init) so gate sharpness can't run away to a - # collapse-inducing extreme during training. + # 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, - "lambda_balance": 0.0, # optional load-balance aux loss weight - # optional entropy-regularization aux loss weight (generic - # Router.entropy_loss, penalizes uniform/collapsed gating) — a - # secondary guard against all experts' widths/temperature - # co-inflating together, which lambda_balance alone can't see - # since per-expert usage shares stay even throughout that - # failure mode. Off by default; bounding above is the primary - # defense. See giant.model.network.Router.entropy_loss. + # 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 - # (see giant.model.network.Router.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 a smooth gradient to every expert. Targets the train/eval - # mismatch identified as a likely contributor to experts - # overlapping instead of partitioning (see CLAUDE.md roadmap). - # gumbel_tau_start/_end are annealed linearly over training - # (giant.train._gumbel_tau); off by default, no effect unless - # gumbel = true. + # 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 - # (0.0 still trains a working router — the gate gets gradient - # through the downstream flow loss like EnergyRouter's centers — - # but only lambda_proc > 0 grounds it in the true `process` label) # 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 (there's no meaningful default axis list) - # — 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 / `--router-axis`. + # 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, + }, } @@ -169,6 +324,9 @@ def auto_device() -> torch.device: # Activation memory is assumed to scale linearly with # batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this # is a rough estimate rather than a guaranteed bound. +# NOTE: not yet recalibrated for the v0.3.0 autoregressive stage-2 trunk — +# see docs/v0.3.0-design.md §9/§11.3, deliberately last in the implementation +# order. _REF_BYTES = 7683 * 1024**2 _REF_BATCH_SIZE = 29696 _REF_HIDDEN_DIM = 1024 @@ -254,13 +412,15 @@ def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None: def load_checkpoint_config(ckpt_path: str | Path) -> dict: - """Load the full ``[train]``/``[model]``/``[meta]`` config.toml written - alongside a checkpoint by ``save_config``. + """Load the full config.toml written alongside a checkpoint by + ``save_config``. Returns ``{}`` if no config.toml sits next to the checkpoint (older runs, or a checkpoint moved without its sidecar) — this is best-effort provenance for threading into a rollout's YAML sidecar, not a hard - requirement for using the checkpoint itself. + requirement for using the checkpoint itself. Returned as-loaded (v0.2 or + v0.3 shape); callers that need the v0.3 shape should run it through + `migrate_config` themselves. """ config_path = Path(ckpt_path).parent / "config.toml" if not config_path.exists(): @@ -283,130 +443,413 @@ def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None: warn_if_git_hash_mismatch(load_toml(config_path), config_path) +def _get_path(d: dict, dotted: str): + """Read a dotted path (e.g. "stage1_model.router.enabled") out of a + nested dict. Returns None if any component along the path is missing.""" + cur = d + for part in dotted.split("."): + if not isinstance(cur, dict) or part not in cur: + return None + cur = cur[part] + return cur + + +def _set_path(d: dict, dotted: str, value) -> None: + """Write a dotted path into a nested dict, creating intermediate dicts as + needed.""" + parts = dotted.split(".") + cur = d + for part in parts[:-1]: + cur = cur.setdefault(part, {}) + cur[parts[-1]] = value + + +def _deep_merge(base: dict, override: dict) -> dict: + """Recursively merge `override` onto a copy of `base`. + + Dict-valued keys recurse instead of being replaced wholesale, so + overriding one leaf (e.g. stage1_model.router.enabled) never drops + untouched siblings — the rest of stage1_model.router, or of + stage1_model — the same property v0.2's router-only bespoke merge had, + generalized here to arbitrary depth. + """ + result = dict(base) + for k, v in override.items(): + if isinstance(v, dict) and isinstance(result.get(k), dict): + result[k] = _deep_merge(result[k], v) + else: + result[k] = v + return result + + +# v0.2 [train] keys that pass through to v0.3 [train] unchanged (same name, +# same meaning) when present in the loaded file — everything model-shaped +# moved to the stage/conditioning blocks instead (see the rest of +# migrate_config below). +_V02_TRAIN_PASSTHROUGH = ( + "epochs", + "batch_size", + "lr", + "weight_decay", + "ema_decay", + "max_val_batches", + "val_fraction", + "num_workers", + "seed", + "validate_every", + "validate_steps", + "warmup_epochs", + "wandb", + "wandb_project", + "wandb_run_name", + "wandb_log_every", +) + +# v0.2 model.hidden_dim/n_blocks/dropout applied identically to both stages +# (there was only ever one trunk shape) -> copied to both stage{1,2}_model. +_V02_MODEL_TO_BOTH_STAGES = ( + ("hidden_dim", "hidden_dim"), + ("n_blocks", "n_res_blocks"), + ("dropout", "dropout"), +) + +# v0.2 train.{n_critic,gp_weight,critic_lr} applied identically to both +# stages' wgan sub-table (there was only ever one wgan objective, shared). +_V02_TRAIN_TO_BOTH_STAGES_WGAN = ( + ("n_critic", "n_critic"), + ("gp_weight", "gp_weight"), + ("critic_lr", "critic_lr"), +) + + +def migrate_config(cfg: dict) -> dict: + """Translate a v0.2 config dict (single [train] + [model]) into the v0.3 + nested format ([conditioning]/[stage1_model]/[stage2_model]/[train]). + + Called on every config.toml load (see merge_cli_overrides) so old + training configs on disk keep working under new code without hand- + editing (decision 3, docs/v0.3.0-design.md §4). `[meta].config_version == + CONFIG_VERSION` marks a dict as already-v0.3; its absence is read as + "this is v0.2" (the design doc's stated rule), so an already-migrated + dict is returned unchanged (deep-copied). + + Only keys actually present in `cfg` are translated — `cfg` may be a + partial file (e.g. `[train]\\nepochs = 5\\n` with no [model] section at + all, relying on v0.2 defaults for everything else). Separately, a fixed + set of v0.2 architectural facts that were never exposed as config keys at + all (e.g. the conditioning MLP was always 2 layers deep, not the v0.3 + default of 1) are injected unconditionally whenever this function decides + it is migrating a v0.2 dict, regardless of which keys the file happened + to set. + + Operates on the config.toml shape. A checkpoint's `model_config` dict + (which additionally carries n_sec_head ownership, §4.1, and needs + `network.build_models`'s cooperation) is a separate migration surface, + deferred to the network.py refactor. + """ + if _get_path(cfg, "meta.config_version") == CONFIG_VERSION: + return copy.deepcopy(cfg) + + cfg = copy.deepcopy(cfg) + old_train = cfg.pop("train", {}) + old_model = cfg.pop("model", {}) + old_router = dict(old_model.pop("router", {})) + + new: dict = {} + + for key in _V02_TRAIN_PASSTHROUGH: + if key in old_train: + _set_path(new, f"train.{key}", old_train[key]) + + if "mode" in old_train: + _set_path(new, "stage1_model.generator", old_train["mode"]) + _set_path(new, "stage2_model.generator", old_train["mode"]) + if "lambda_nsec" in old_train: + _set_path(new, "stage2_model.n_sec.lambda", old_train["lambda_nsec"]) + if "lambda_s2" in old_train: + _set_path(new, "stage2_model.lambda", old_train["lambda_s2"]) + for old_key, new_key in _V02_TRAIN_TO_BOTH_STAGES_WGAN: + if old_key in old_train: + _set_path(new, f"stage1_model.wgan.{new_key}", old_train[old_key]) + _set_path(new, f"stage2_model.wgan.{new_key}", old_train[old_key]) + + for old_key, new_key in _V02_MODEL_TO_BOTH_STAGES: + if old_key in old_model: + _set_path(new, f"stage1_model.{new_key}", old_model[old_key]) + _set_path(new, f"stage2_model.{new_key}", old_model[old_key]) + if "emb_dim" in old_model: + _set_path(new, "conditioning.particle.emb_dim", old_model["emb_dim"]) + _set_path(new, "conditioning.material.emb_dim", old_model["emb_dim"]) + if "conditioning" in old_model: + _set_path(new, "conditioning.particle.type", old_model["conditioning"]) + _set_path(new, "conditioning.material.type", old_model["conditioning"]) + if "noise_dim" in old_model: + _set_path(new, "stage1_model.wgan.noise_dim", old_model["noise_dim"]) + _set_path(new, "stage2_model.wgan.noise_dim", old_model["noise_dim"]) + if "k_max" in old_model: + _set_path(new, "stage2_model.k_max", old_model["k_max"]) + + if old_router: + expert_hidden_dim = old_router.pop("expert_hidden_dim", 0) + expert_n_blocks = old_router.pop("expert_n_blocks", 0) + if expert_hidden_dim or expert_n_blocks: + raise ValueError( + "v0.2 config sets model.router.expert_hidden_dim/" + f"expert_n_blocks to a non-default value " + f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed " + "per-expert sizing (experts always inherit the stage's " + "hidden_dim/n_res_blocks), so this config's routed experts " + "have a different width/depth than the monolith and its " + "checkpoint can only be loaded by v0.2 code." + ) + _set_path(new, "stage1_model.router", dict(old_router)) + stage2_router = dict(old_router) + stage2_router["tie_to_stage1"] = False + _set_path(new, "stage2_model.router", stage2_router) + + # v0.2 architectural facts with no corresponding config key at all — + # always set once we've determined we're migrating a v0.2 dict, + # independent of what the file did/didn't specify. NOTE: n_layers here + # (2) differs from the v0.3 *default* (1) — this is not a typo, see the + # docstring above. + _set_path(new, "conditioning.out_dim", 128) + _set_path(new, "conditioning.particle.n_layers", 2) + _set_path(new, "conditioning.material.n_layers", 2) + _set_path(new, "stage1_model.active", True) + _set_path(new, "stage1_model.flow.time_dim", 64) + _set_path(new, "stage1_model.ddpm.time_dim", 64) + _set_path(new, "stage2_model.active", True) + _set_path(new, "stage2_model.flow.time_dim", 64) + _set_path(new, "stage2_model.ddpm.time_dim", 64) + _set_path(new, "stage2_model.context_dim", 64) + _set_path(new, "stage2_model.decoder", "one_shot") + _set_path(new, "stage2_model.particle_type.target", "physical") + + new_meta = dict(cfg.pop("meta", {})) + new_meta["config_version"] = CONFIG_VERSION + new["meta"] = new_meta + + # Anything else in the original dict (unrecognized top-level sections) + # carries through untouched rather than being silently dropped. + for k, v in cfg.items(): + new.setdefault(k, v) + + return new + + def merge_cli_overrides( defaults: dict, config_path: Path | None, - train_overrides: dict, - model_overrides: dict, + overrides: dict, ) -> dict: - """Resolve config as defaults -> TOML file -> explicit CLI flags. + """Resolve config as defaults -> TOML file -> explicit overrides. - `model.router` is deep-merged one level (rather than replaced wholesale) - at each stage, so a TOML file or CLI flag only overriding e.g. - `router.enabled` doesn't drop the rest of the router defaults. + `overrides` is keyed by top-level section name (e.g. "stage1_model", + "train"), each value an arbitrarily nested dict of overrides to + deep-merge (see `_deep_merge`) — the shape stage-prefixed CLI flags + naturally produce. A v0.2-shaped TOML file is transparently migrated + (`migrate_config`) before merging, so old configs on disk keep working + under the new schema. """ - cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])} - cfg["model"]["router"] = dict(defaults["model"]["router"]) + cfg = copy.deepcopy(defaults) if config_path is not None: - file_cfg = load_toml(config_path) - cfg["train"].update(file_cfg.get("train", {})) - file_model = dict(file_cfg.get("model", {})) - file_router = file_model.pop("router", None) - cfg["model"].update(file_model) - if file_router: - cfg["model"]["router"].update(file_router) + file_cfg = migrate_config(load_toml(config_path)) + for section, values in file_cfg.items(): + if section == "meta": + continue + if isinstance(values, dict): + cfg[section] = _deep_merge(cfg.get(section, {}), values) + else: + cfg[section] = values warn_if_git_hash_mismatch(file_cfg, config_path) - model_overrides = dict(model_overrides) - router_overrides = model_overrides.pop("router", None) - cfg["train"].update(train_overrides) - cfg["model"].update(model_overrides) - if router_overrides: - cfg["model"]["router"].update(router_overrides) + for section, values in overrides.items(): + if isinstance(values, dict): + cfg[section] = _deep_merge(cfg.get(section, {}), values) + else: + cfg[section] = values return cfg -def resolve_expert_dims( - router_cfg: dict, hidden_dim: int, n_blocks: int -) -> tuple[int, int]: - """Resolve a router's expert hidden_dim/n_blocks, inheriting from the - monolith's when left at the 0 ("unset") sentinel. +def validate_config(cfg: dict) -> None: + """Cross-block validation the per-block schema can't express on its own. - Used by both `giant.pipeline` (to build the checkpoint's `model_config`) - and `giant.cli`'s batch-size auto-estimate, so `--hidden-dim`/`--n-blocks` - size the experts the same way in both places unless - `router.expert_hidden_dim`/`expert_n_blocks` are explicitly overridden. + Raises ValueError with a clear message on the first violation found. Call + after `merge_cli_overrides` has produced a fully-merged v0.3 config — + these checks need to see across blocks, so they don't belong in + `migrate_config` (which only ever sees one dict's own keys) or in any + single block's defaults. """ - expert_hidden_dim = router_cfg.get("expert_hidden_dim") or hidden_dim - expert_n_blocks = router_cfg.get("expert_n_blocks") or n_blocks - return expert_hidden_dim, expert_n_blocks + particle_type = _get_path(cfg, "conditioning.particle.type") + + pt_target = _get_path(cfg, "stage2_model.particle_type.target") + if pt_target == "embedding" and particle_type != "embedding": + raise ValueError( + "stage2_model.particle_type.target = 'embedding' requires " + "conditioning.particle.type = 'embedding' (there is no embedding " + "table to regress against under conditioning.particle.type = " + f"{particle_type!r})" + ) + + for stage_name in ("stage1_model", "stage2_model"): + router = _get_path(cfg, f"{stage_name}.router") or {} + if ( + router.get("enabled") + and router.get("type") in ("pdg", "process") + and particle_type == "physical" + ): + raise ValueError( + f"{stage_name}.router.type = {router['type']!r} builds its " + "own training-vocab-scoped embedding, incompatible with " + "conditioning.particle.type = 'physical' (defeats " + "generalization beyond the training menu) — pick a " + "different router type or a different " + "conditioning.particle.type" + ) + + if _get_path(cfg, "stage2_model.router.tie_to_stage1") and not _get_path( + cfg, "stage1_model.active" + ): + raise ValueError( + "stage2_model.router.tie_to_stage1 = true requires " + "stage1_model.active = true (there is no stage-1 router to tie to)" + ) + + if _get_path(cfg, "stage2_model.n_sec.mode") == "stop_token": + raise ValueError( + "stage2_model.n_sec.mode = 'stop_token' is accepted by the schema " + "but not implemented in v0.3.0 — use 'head' (default) or 'truth' " + "(standalone stage-2 evaluation only, never for rollout)" + ) -_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb"} +_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb", "onehot": "oh"} -# Priority-ordered candidate fields for default_out_dir_name: (label, getter, -# formatter). `getter(train, model)` returns None when the field is at its -# default (and so should be omitted); otherwise formatter(value) renders the -# name token. The router is a single unit gated on `router.enabled` rather -# than one candidate per router key, since its type/n_experts are meaningless -# while disabled. -def _mode_candidate(train, model): - return None if train["mode"] == DEFAULT_CONFIG["train"]["mode"] else train["mode"] +def _path_candidate(dotted_path: str, prefix: str, formatter=str): + """Candidate factory: show `prefix + formatter(value)` when the value at + `dotted_path` differs from its DEFAULT_CONFIG value, else omit.""" - -def _router_candidate(train, model): - router = model["router"] - if router["enabled"] == DEFAULT_CONFIG["model"]["router"]["enabled"]: - return None - return f"r-{router['type']}{router['n_experts']}" - - -def _router_flag_candidate(field, token_map): - """Candidate factory for a boolean `model.router` sub-field. - - Gated on `router.enabled` like `_router_candidate` (a disabled router's - sub-fields are meaningless), then omitted unless `field` differs from - its DEFAULT_CONFIG value — same "only show non-default" rule as every - other candidate. `token_map` need only cover the non-default value(s), - since the default value always yields None. - """ - - def _candidate(train, model): - router = model["router"] - default_router = DEFAULT_CONFIG["model"]["router"] - if router["enabled"] == default_router["enabled"]: + def _candidate(cfg): + value = _get_path(cfg, dotted_path) + default = _get_path(DEFAULT_CONFIG, dotted_path) + if value == default: return None - value = router[field] - if value == default_router[field]: - return None - return token_map[value] + return f"{prefix}{formatter(value)}" return _candidate -def _conditioning_candidate(train, model): - if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]: - return None - code = _CONDITIONING_CODE.get(model["conditioning"], model["conditioning"]) - return f"c{code}" - - -def _default_field_candidate(section_key, field, prefix): - def _candidate(train, model): - section = train if section_key == "train" else model - value = section[field] - if value == DEFAULT_CONFIG[section_key][field]: +def _conditioning_candidate(axis: str, short: str): + def _candidate(cfg): + value = _get_path(cfg, f"conditioning.{axis}.type") + default = _get_path(DEFAULT_CONFIG, f"conditioning.{axis}.type") + if value == default: return None - return f"{prefix}{value}" + code = _CONDITIONING_CODE.get(value, value) + return f"{short}{code}" return _candidate +def _router_candidate(stage_key: str, short: str): + """Candidate for a stage's router as a single unit, gated on + `router.enabled` (a disabled router's type/n_experts are meaningless).""" + + def _candidate(cfg): + router = _get_path(cfg, f"{stage_key}.router") or {} + default_router = _get_path(DEFAULT_CONFIG, f"{stage_key}.router") or {} + if router.get("enabled") == default_router.get("enabled"): + return None + return f"{short}r-{router['type']}{router['n_experts']}" + + return _candidate + + +def _router_flag_candidate(stage_key: str, short: str, field: str, token_map: dict): + """Candidate factory for a boolean field inside a stage's router block. + + Gated on `router.enabled` like `_router_candidate`, then omitted unless + `field` differs from its DEFAULT_CONFIG value. `token_map` need only + cover the non-default value(s), since the default value always yields + None. + """ + + def _candidate(cfg): + router = _get_path(cfg, f"{stage_key}.router") or {} + default_router = _get_path(DEFAULT_CONFIG, f"{stage_key}.router") or {} + if router.get("enabled") == default_router.get("enabled"): + return None + value = router.get(field) + if value == default_router.get(field): + return None + return f"{short}{token_map[value]}" + + return _candidate + + +# Priority-ordered candidate fields for default_out_dir_name: (label, +# candidate(cfg) -> str | None). Beyond _OUT_DIR_NAME_MAX_FIELDS non-default +# fields, the remainder collapse into a hash suffix (see +# default_out_dir_name). Candidates read the whole nested cfg via dotted +# paths — there is no single "model" dict anymore now that architecture is +# split across conditioning/stage1_model/stage2_model. _OUT_DIR_NAME_CANDIDATES = [ - ("mode", _mode_candidate), - ("router", _router_candidate), - ("gumbel", _router_flag_candidate("gumbel", {True: "gum"})), - ("learn_centers", _router_flag_candidate("learn_centers", {False: "nolc"})), - ("learn_width", _router_flag_candidate("learn_width", {True: "lw"})), - ("learn_temperature", _router_flag_candidate("learn_temperature", {True: "lt"})), - ("conditioning", _conditioning_candidate), - ("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")), - ("n_blocks", _default_field_candidate("model", "n_blocks", "b")), - ("emb_dim", _default_field_candidate("model", "emb_dim", "e")), - ("lr", _default_field_candidate("train", "lr", "lr")), - ("batch_size", _default_field_candidate("train", "batch_size", "bs")), - ("seed", _default_field_candidate("train", "seed", "seed")), - ("epochs", _default_field_candidate("train", "epochs", "ep")), + ("stage1_generator", _path_candidate("stage1_model.generator", "")), + ("stage2_generator", _path_candidate("stage2_model.generator", "s2-")), + ("stage2_decoder", _path_candidate("stage2_model.decoder", "dec-")), + ( + "stage2_history", + _path_candidate("stage2_model.autoregressive.history", "hist-"), + ), + ( + "particle_type_target", + _path_candidate("stage2_model.particle_type.target", "pt-"), + ), + ("stage1_router", _router_candidate("stage1_model", "s1")), + ("stage2_router", _router_candidate("stage2_model", "s2")), + ( + "stage1_gumbel", + _router_flag_candidate("stage1_model", "s1", "gumbel", {True: "gum"}), + ), + ( + "stage2_gumbel", + _router_flag_candidate("stage2_model", "s2", "gumbel", {True: "gum"}), + ), + ( + "stage1_learn_centers", + _router_flag_candidate("stage1_model", "s1", "learn_centers", {False: "nolc"}), + ), + ( + "stage2_learn_centers", + _router_flag_candidate("stage2_model", "s2", "learn_centers", {False: "nolc"}), + ), + ( + "stage1_learn_width", + _router_flag_candidate("stage1_model", "s1", "learn_width", {True: "lw"}), + ), + ( + "stage2_learn_width", + _router_flag_candidate("stage2_model", "s2", "learn_width", {True: "lw"}), + ), + ( + "stage1_learn_temperature", + _router_flag_candidate("stage1_model", "s1", "learn_temperature", {True: "lt"}), + ), + ( + "stage2_learn_temperature", + _router_flag_candidate("stage2_model", "s2", "learn_temperature", {True: "lt"}), + ), + ("particle_conditioning", _conditioning_candidate("particle", "c")), + ("material_conditioning", _conditioning_candidate("material", "m")), + ("stage1_hidden_dim", _path_candidate("stage1_model.hidden_dim", "h")), + ("stage2_hidden_dim", _path_candidate("stage2_model.hidden_dim", "s2h")), + ("stage1_n_res_blocks", _path_candidate("stage1_model.n_res_blocks", "b")), + ("stage2_n_res_blocks", _path_candidate("stage2_model.n_res_blocks", "s2b")), + ("particle_emb_dim", _path_candidate("conditioning.particle.emb_dim", "e")), + ("lr", _path_candidate("train.lr", "lr")), + ("batch_size", _path_candidate("train.batch_size", "bs")), + ("seed", _path_candidate("train.seed", "seed")), + ("epochs", _path_candidate("train.epochs", "ep")), ] _OUT_DIR_NAME_MAX_FIELDS = 6 @@ -424,11 +867,10 @@ def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str: giant.train), which is the reason a timestamp is always included. """ now = now or datetime.now() - train, model = cfg["train"], cfg["model"] tokens = [] overflow = [] for label, candidate in _OUT_DIR_NAME_CANDIDATES: - token = candidate(train, model) + token = candidate(cfg) if token is None: continue if len(tokens) < _OUT_DIR_NAME_MAX_FIELDS: @@ -476,30 +918,31 @@ def _toml_value(v) -> str: return str(v) -def save_config(cfg: dict, out_dir: Path, meta: dict) -> None: - lines = [] - # One-level-nested dict values (e.g. model.router) are rendered as their - # own [section.subsection] table after the parent section, since TOML - # doesn't accept a bare dict as a `key = value` scalar line. - nested_sections: list[tuple[str, dict]] = [] - for section, values in cfg.items(): - lines.append(f"[{section}]") - for k, v in values.items(): - if isinstance(v, dict): - nested_sections.append((f"{section}.{k}", v)) - continue - lines.append(f"{k:<14} = {_toml_value(v)}") - lines.append("") +def _write_section(lines: list[str], path: str, values: dict) -> None: + """Write one TOML table (`[path]`) and recurse depth-first into any + dict-valued keys as `[path.subkey]` — handles the v0.3 schema's 2-3 level + nesting (e.g. stage1_model.router, stage2_model.n_sec) with no depth + limit, unlike the one-level-only writer this replaces.""" + lines.append(f"[{path}]") + nested: list[tuple[str, dict]] = [] + for k, v in values.items(): + if isinstance(v, dict): + nested.append((f"{path}.{k}", v)) + else: + lines.append(f"{k:<18} = {_toml_value(v)}") + lines.append("") + for sub_path, sub_values in nested: + _write_section(lines, sub_path, sub_values) - for name, values in nested_sections: - lines.append(f"[{name}]") - for k, v in values.items(): - lines.append(f"{k:<14} = {_toml_value(v)}") - lines.append("") + +def save_config(cfg: dict, out_dir: Path, meta: dict) -> None: + lines: list[str] = [] + for section, values in cfg.items(): + _write_section(lines, section, values) lines.append("[meta]") for k, v in meta.items(): - lines.append(f"{k:<14} = {_toml_value(v)}") + lines.append(f"{k:<18} = {_toml_value(v)}") (out_dir / "config.toml").write_text("\n".join(lines)) @@ -514,6 +957,7 @@ def build_run_meta( n_train_steps: int, ) -> dict: return { + "config_version": CONFIG_VERSION, "git_hash": git_hash(), "seed": seed, "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"), diff --git a/tests/test_config.py b/tests/test_config.py index 30f6e41..eee0f18 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,46 +1,314 @@ from datetime import datetime +from pathlib import Path from giant import config as gconfig +_CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs" -def _write_config(path, git_hash): - path.write_text( - f""" -[train] -epochs = 5 -[model] -hidden_dim = 64 +def _write_toml(path, git_hash=None, extra=""): + meta = f'\n[meta]\ngit_hash = "{git_hash}"\n' if git_hash is not None else "" + path.write_text(extra + meta) -[meta] -git_hash = "{git_hash}" -""" + +# --------------------------------------------------------------------------- +# Conditioning enum +# --------------------------------------------------------------------------- + + +def test_conditioning_enum_has_onehot(): + assert gconfig.Conditioning.onehot == "onehot" + assert {c.value for c in gconfig.Conditioning} == { + "physical", + "embedding", + "onehot", + } + + +# --------------------------------------------------------------------------- +# _deep_merge +# --------------------------------------------------------------------------- + + +def test_deep_merge_leaf_override_keeps_untouched_siblings(): + base = {"a": 1, "b": {"c": 2, "d": 3}} + result = gconfig._deep_merge(base, {"b": {"c": 99}}) + assert result == {"a": 1, "b": {"c": 99, "d": 3}} + + +def test_deep_merge_recurses_at_multiple_levels(): + base = { + "stage1_model": { + "hidden_dim": 256, + "router": {"enabled": False, "type": "energy", "n_experts": 4}, + } + } + result = gconfig._deep_merge(base, {"stage1_model": {"router": {"enabled": True}}}) + assert result["stage1_model"]["hidden_dim"] == 256 + assert result["stage1_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 4, + } + + +def test_deep_merge_does_not_mutate_base(): + base = {"a": {"b": 1}} + gconfig._deep_merge(base, {"a": {"b": 2}}) + assert base == {"a": {"b": 1}} + + +def test_deep_merge_non_dict_override_replaces_wholesale(): + base = {"a": {"b": 1}} + result = gconfig._deep_merge(base, {"a": 5}) + assert result == {"a": 5} + + +# --------------------------------------------------------------------------- +# migrate_config +# --------------------------------------------------------------------------- + + +def test_migrate_config_already_v3_returned_unchanged(): + cfg = {"meta": {"config_version": 3}, "stage1_model": {"generator": "flow"}} + result = gconfig.migrate_config(cfg) + assert result == cfg + result["stage1_model"]["generator"] = "wgan" + assert cfg["stage1_model"]["generator"] == "flow" # deep-copied, not aliased + + +def test_migrate_config_empty_dict_still_injects_hardcoded_v02_facts(): + # No [train]/[model] at all still counts as "v0.2" (config_version + # absent) — the hardcoded architectural facts fire unconditionally. + new = gconfig.migrate_config({}) + assert "train" not in new + assert new["conditioning"]["out_dim"] == 128 + assert new["conditioning"]["particle"]["n_layers"] == 2 + assert new["conditioning"]["material"]["n_layers"] == 2 + assert new["stage1_model"]["active"] is True + assert new["stage1_model"]["flow"]["time_dim"] == 64 + assert new["stage1_model"]["ddpm"]["time_dim"] == 64 + assert new["stage2_model"]["active"] is True + assert new["stage2_model"]["flow"]["time_dim"] == 64 + assert new["stage2_model"]["ddpm"]["time_dim"] == 64 + assert new["stage2_model"]["context_dim"] == 64 + assert new["stage2_model"]["decoder"] == "one_shot" + assert new["stage2_model"]["particle_type"]["target"] == "physical" + assert new["meta"] == {"config_version": 3} + + +def test_migrate_config_mode_maps_to_both_stage_generators(): + new = gconfig.migrate_config({"train": {"mode": "wgan"}}) + assert new["stage1_model"]["generator"] == "wgan" + assert new["stage2_model"]["generator"] == "wgan" + + +def test_migrate_config_lambda_nsec_and_lambda_s2(): + new = gconfig.migrate_config({"train": {"lambda_nsec": 0.2, "lambda_s2": 2.0}}) + assert new["stage2_model"]["n_sec"]["lambda"] == 0.2 + assert new["stage2_model"]["lambda"] == 2.0 + + +def test_migrate_config_wgan_knobs_map_to_both_stages(): + new = gconfig.migrate_config( + {"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}} ) + for stage in ("stage1_model", "stage2_model"): + assert new[stage]["wgan"]["n_critic"] == 3 + assert new[stage]["wgan"]["gp_weight"] == 5.0 + assert new[stage]["wgan"]["critic_lr"] == 1e-4 -def test_merge_cli_overrides_applies_file_then_cli(tmp_path, monkeypatch): +def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages(): + new = gconfig.migrate_config( + {"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}} + ) + for stage in ("stage1_model", "stage2_model"): + assert new[stage]["hidden_dim"] == 128 + assert new[stage]["n_res_blocks"] == 4 + assert new[stage]["dropout"] == 0.2 + + +def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes(): + new = gconfig.migrate_config( + {"model": {"emb_dim": 32, "conditioning": "embedding"}} + ) + for axis in ("particle", "material"): + assert new["conditioning"][axis]["emb_dim"] == 32 + assert new["conditioning"][axis]["type"] == "embedding" + + +def test_migrate_config_noise_dim_maps_to_both_stages_wgan(): + new = gconfig.migrate_config({"model": {"noise_dim": 128}}) + assert new["stage1_model"]["wgan"]["noise_dim"] == 128 + assert new["stage2_model"]["wgan"]["noise_dim"] == 128 + + +def test_migrate_config_k_max_maps_to_stage2_only(): + new = gconfig.migrate_config({"model": {"k_max": 20}}) + assert new["stage2_model"]["k_max"] == 20 + assert "k_max" not in new.get("stage1_model", {}) + + +def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false(): + new = gconfig.migrate_config( + { + "model": { + "router": { + "enabled": True, + "type": "energy", + "n_experts": 10, + "temperature": 0.05, + } + } + } + ) + assert new["stage1_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 10, + "temperature": 0.05, + } + assert new["stage2_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 10, + "temperature": 0.05, + "tie_to_stage1": False, + } + + +def test_migrate_config_router_nonzero_expert_dims_raises(): + cfg = { + "model": { + "router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0} + } + } + try: + gconfig.migrate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "expert_hidden_dim" in str(e) + + +def test_migrate_config_router_zero_expert_dims_dropped_silently(): + new = gconfig.migrate_config( + { + "model": { + "router": { + "enabled": True, + "type": "energy", + "n_experts": 4, + "expert_hidden_dim": 0, + "expert_n_blocks": 0, + } + } + } + ) + assert "expert_hidden_dim" not in new["stage1_model"]["router"] + assert "expert_n_blocks" not in new["stage1_model"]["router"] + + +def test_migrate_config_train_passthrough_is_exact(): + new = gconfig.migrate_config({"train": {"epochs": 7, "batch_size": 999, "seed": 3}}) + assert new["train"] == {"epochs": 7, "batch_size": 999, "seed": 3} + + +def test_migrate_config_preserves_meta_git_hash(): + new = gconfig.migrate_config({"meta": {"git_hash": "abc123"}}) + assert new["meta"] == {"git_hash": "abc123", "config_version": 3} + + +def test_migrate_config_real_router_fixture_raises_on_nonzero_expert_dims(): + cfg = gconfig.load_toml(_CONFIGS_DIR / "router_energy_n10_embedding.toml") + try: + gconfig.migrate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "expert_hidden_dim" in str(e) + + +def test_migrate_config_real_wgan_fixture(): + cfg = gconfig.load_toml(_CONFIGS_DIR / "wgan_h128_b4_physical.toml") + new = gconfig.migrate_config(cfg) + assert new["stage1_model"]["generator"] == "wgan" + assert new["stage2_model"]["generator"] == "wgan" + for stage in ("stage1_model", "stage2_model"): + assert new[stage]["hidden_dim"] == 128 + assert new[stage]["n_res_blocks"] == 4 + assert new[stage]["dropout"] == 0.0 + assert new["conditioning"]["particle"]["type"] == "physical" + assert new["conditioning"]["material"]["type"] == "physical" + assert new["train"]["epochs"] == 30 + assert new["train"]["warmup_epochs"] == 3 + + +# --------------------------------------------------------------------------- +# merge_cli_overrides +# --------------------------------------------------------------------------- + + +def test_merge_cli_overrides_defaults_only_matches_default_config(): + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}) + assert cfg == gconfig.DEFAULT_CONFIG + assert cfg is not gconfig.DEFAULT_CONFIG + + +def test_merge_cli_overrides_nested_override_keeps_siblings(): + cfg = gconfig.merge_cli_overrides( + gconfig.DEFAULT_CONFIG, + None, + {"stage1_model": {"router": {"enabled": True}}}, + ) + assert cfg["stage1_model"]["router"]["enabled"] is True + assert cfg["stage1_model"]["router"]["type"] == "energy" # default preserved + assert cfg["stage1_model"]["hidden_dim"] == 256 # untouched sibling section + + +def test_merge_cli_overrides_file_then_explicit_override_precedence( + tmp_path, monkeypatch +): monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123") path = tmp_path / "config.toml" - _write_config(path, "abc123") + _write_toml( + path, + git_hash="abc123", + extra="[train]\nepochs = 5\n\n[model]\nhidden_dim = 64\n", + ) cfg = gconfig.merge_cli_overrides( gconfig.DEFAULT_CONFIG, path, - train_overrides={}, - model_overrides={"hidden_dim": 128}, + {"stage1_model": {"hidden_dim": 128}}, ) assert cfg["train"]["epochs"] == 5 # from file - assert cfg["model"]["hidden_dim"] == 128 # CLI override wins over file + assert cfg["stage1_model"]["hidden_dim"] == 128 # explicit override wins over file + assert cfg["stage2_model"]["hidden_dim"] == 64 # migrated from file, not overridden + + +def test_merge_cli_overrides_migrates_v2_file_transparently(tmp_path, monkeypatch): + monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123") + path = tmp_path / "config.toml" + _write_toml( + path, + git_hash="abc123", + extra='[train]\nmode = "wgan"\n\n[model]\nconditioning = "embedding"\n', + ) + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) + assert cfg["stage1_model"]["generator"] == "wgan" + assert cfg["stage2_model"]["generator"] == "wgan" + assert cfg["conditioning"]["particle"]["type"] == "embedding" + # hardcoded v0.2 fact still applied even though it's not a CLI-settable key + assert cfg["conditioning"]["particle"]["n_layers"] == 2 def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") path = tmp_path / "config.toml" - _write_config(path, "old111") + _write_toml(path, git_hash="old111", extra="[train]\nepochs = 5\n") - cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) - - assert cfg["train"]["epochs"] == 5 # does not fail, config still applied + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) captured = capsys.readouterr() assert "warning" in captured.err assert "old111" in captured.err @@ -52,9 +320,9 @@ def test_merge_cli_overrides_no_warning_on_matching_git_hash( ): monkeypatch.setattr(gconfig, "git_hash", lambda: "same123") path = tmp_path / "config.toml" - _write_config(path, "same123") + _write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n") - gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" @@ -63,9 +331,9 @@ def test_merge_cli_overrides_no_warning_when_git_hash_unknown( ): monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown") path = tmp_path / "config.toml" - _write_config(path, "abc123") + _write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n") - gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" @@ -76,17 +344,295 @@ def test_merge_cli_overrides_no_warning_when_meta_section_absent( path = tmp_path / "config.toml" path.write_text("[train]\nepochs = 5\n") - gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" +def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch): + monkeypatch.setattr( + gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243" + ) + cfg = gconfig.merge_cli_overrides( + gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {} + ) + assert cfg["stage1_model"]["generator"] == "flow" + assert cfg["stage1_model"]["hidden_dim"] == 256 + assert cfg["stage2_model"]["hidden_dim"] == 256 + assert cfg["conditioning"]["particle"]["emb_dim"] == 16 + assert cfg["conditioning"]["particle"]["n_layers"] == 2 # migrated hardcoded fact + assert cfg["train"]["epochs"] == 100 + assert cfg["stage2_model"]["decoder"] == "one_shot" + + +# --------------------------------------------------------------------------- +# save_config +# --------------------------------------------------------------------------- + + +def test_save_config_round_trips_multi_level_nesting(tmp_path): + cfg = { + "stage1_model": { + "hidden_dim": 256, + "router": { + "enabled": True, + "type": "energy", + "n_experts": 4, + }, + }, + "train": {"epochs": 100}, + } + meta = {"config_version": 3, "git_hash": "abc123"} + + gconfig.save_config(cfg, tmp_path, meta) + loaded = gconfig.load_toml(tmp_path / "config.toml") + + assert loaded["stage1_model"]["hidden_dim"] == 256 + assert loaded["stage1_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 4, + } + assert loaded["train"] == {"epochs": 100} + assert loaded["meta"] == meta + + +def test_save_config_round_trips_three_level_nesting(tmp_path): + cfg = { + "stage2_model": { + "decoder": "autoregressive", + "n_sec": {"mode": "head", "lambda": 0.1}, + "router": {"tie_to_stage1": True}, + } + } + gconfig.save_config(cfg, tmp_path, {"config_version": 3}) + loaded = gconfig.load_toml(tmp_path / "config.toml") + + assert loaded["stage2_model"]["decoder"] == "autoregressive" + assert loaded["stage2_model"]["n_sec"] == {"mode": "head", "lambda": 0.1} + assert loaded["stage2_model"]["router"] == {"tie_to_stage1": True} + + +# --------------------------------------------------------------------------- +# default_out_dir_name +# --------------------------------------------------------------------------- + +_NOW = datetime(2026, 7, 29, 14, 30) + + +def _cfg_with(**dotted_overrides): + """Build a full DEFAULT_CONFIG-shaped dict with dotted-path overrides + applied via _deep_merge, e.g. _cfg_with(**{"stage1_model.hidden_dim": 512}).""" + overrides: dict = {} + for dotted, value in dotted_overrides.items(): + gconfig._set_path(overrides, dotted, value) + return gconfig._deep_merge(gconfig.DEFAULT_CONFIG, overrides) + + +def test_default_out_dir_name_all_defaults_is_just_the_timestamp(): + assert ( + gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) + == "20260729_1430" + ) + + +def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix(): + cfg = _cfg_with(**{"stage1_model.generator": "wgan"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan" + + +def test_default_out_dir_name_stage2_decoder_shown(): + cfg = _cfg_with(**{"stage2_model.decoder": "one_shot"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_dec-one_shot" + + +def test_default_out_dir_name_particle_type_target_shown(): + cfg = _cfg_with(**{"stage2_model.particle_type.target": "physical"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_pt-physical" + + +def test_default_out_dir_name_particle_conditioning_embedding_abbreviated(): + cfg = _cfg_with(**{"conditioning.particle.type": "embedding"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb" + + +def test_default_out_dir_name_stage1_router_shown_as_unit(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + } + ) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8" + + +def test_default_out_dir_name_stage2_router_shown_as_unit_distinct_from_stage1(): + cfg = _cfg_with( + **{ + "stage2_model.router.enabled": True, + "stage2_model.router.type": "pdg", + "stage2_model.router.n_experts": 3, + } + ) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s2r-pdg3" + + +def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": False, + "stage1_model.router.type": "pdg", + "stage1_model.router.n_experts": 8, + } + ) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" + + +def test_default_out_dir_name_router_gumbel_shown_when_enabled(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + "stage1_model.router.gumbel": True, + } + ) + assert ( + gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum" + ) + + +def test_default_out_dir_name_overflow_caps_and_hashes_remainder(): + cfg = _cfg_with( + **{ + "stage1_model.generator": "wgan", + "stage2_model.generator": "flow", + "stage2_model.decoder": "one_shot", + "stage2_model.autoregressive.history": "attention", + "stage2_model.particle_type.target": "physical", + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + "stage2_model.router.enabled": True, + "stage2_model.router.type": "pdg", + "stage2_model.router.n_experts": 3, + } + ) + name = gconfig.default_out_dir_name(cfg, now=_NOW) + # First 6 by priority: stage1_generator, stage2_generator, stage2_decoder, + # stage2_history, particle_type_target, stage1_router — stage2_router + # overflows into the hash suffix. + assert name.startswith( + "20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+" + ) + + +def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive(): + overrides = { + "stage1_model.generator": "wgan", + "stage2_model.generator": "flow", + "stage2_model.decoder": "one_shot", + "stage2_model.autoregressive.history": "attention", + "stage2_model.particle_type.target": "physical", + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + "train.seed": 3, + } + name_a = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW) + name_b = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW) + assert name_a == name_b + + changed = dict(overrides, **{"train.seed": 99}) + name_c = gconfig.default_out_dir_name(_cfg_with(**changed), now=_NOW) + assert name_c != name_a + + +# --------------------------------------------------------------------------- +# validate_config +# --------------------------------------------------------------------------- + + +def test_validate_config_default_config_passes(): + gconfig.validate_config(gconfig.DEFAULT_CONFIG) # must not raise + + +def test_validate_config_embedding_target_requires_embedding_conditioning(): + cfg = _cfg_with( + **{ + "stage2_model.particle_type.target": "embedding", + "conditioning.particle.type": "physical", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "embedding" in str(e) + + +def test_validate_config_embedding_target_passes_with_embedding_conditioning(): + cfg = _cfg_with( + **{ + "stage2_model.particle_type.target": "embedding", + "conditioning.particle.type": "embedding", + } + ) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_pdg_router_incompatible_with_physical_conditioning(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": True, + "stage1_model.router.type": "pdg", + "conditioning.particle.type": "physical", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "pdg" in str(e) + + +def test_validate_config_tie_to_stage1_requires_stage1_active(): + cfg = _cfg_with( + **{ + "stage2_model.router.tie_to_stage1": True, + "stage1_model.active": False, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "tie_to_stage1" in str(e) + + +def test_validate_config_stop_token_not_implemented(): + cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"}) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "stop_token" in str(e) + + +# --------------------------------------------------------------------------- +# checkpoint config-mismatch warnings (unchanged surface, still exercised) +# --------------------------------------------------------------------------- + + def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml( tmp_path, monkeypatch, capsys ): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") ckpt_path = tmp_path / "best.pt" - ckpt_path.write_bytes(b"") # contents irrelevant, only its directory is used - _write_config(tmp_path / "config.toml", "old111") + ckpt_path.write_bytes(b"") + _write_toml( + tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n" + ) gconfig.warn_if_checkpoint_config_mismatch(ckpt_path) @@ -113,202 +659,9 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match( monkeypatch.setattr(gconfig, "git_hash", lambda: "same123") ckpt_path = tmp_path / "best.pt" ckpt_path.write_bytes(b"") - _write_config(tmp_path / "config.toml", "same123") + _write_toml( + tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n" + ) gconfig.warn_if_checkpoint_config_mismatch(ckpt_path) assert capsys.readouterr().err == "" - - -def test_resolve_expert_dims_default_config_inherits_hidden_dim_and_n_blocks(): - # The unset sentinel (expert_hidden_dim/n_blocks == 0 in DEFAULT_CONFIG) - # is exactly the bug fixed by resolve_expert_dims: it must not silently - # fall back to some other hardcoded default, only to the monolith's own - # hidden_dim/n_blocks, so --hidden-dim/--n-blocks reach the experts too. - router_cfg = dict(gconfig.DEFAULT_CONFIG["model"]["router"]) - assert router_cfg["expert_hidden_dim"] == 0 - assert router_cfg["expert_n_blocks"] == 0 - - hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6) - assert (hidden_dim, n_blocks) == (512, 6) - - -def test_resolve_expert_dims_missing_keys_also_inherit(): - hidden_dim, n_blocks = gconfig.resolve_expert_dims({}, 512, 6) - assert (hidden_dim, n_blocks) == (512, 6) - - -def test_default_config_gumbel_router_defaults_off(): - # Straight-through Gumbel-softmax combine weights (giant.model.network. - # Router.combine_weights) must be opt-in — existing routed configs and - # checkpoints should be unaffected unless gumbel is explicitly enabled. - router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"] - assert router_cfg["gumbel"] is False - assert router_cfg["gumbel_tau_start"] == 1.0 - assert router_cfg["gumbel_tau_end"] == 0.1 - - -def test_resolve_expert_dims_explicit_override_wins(): - router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3} - hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6) - assert (hidden_dim, n_blocks) == (128, 3) - - -def test_resolve_expert_dims_partial_override_mixes_explicit_and_inherited(): - router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 0} - hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6) - assert (hidden_dim, n_blocks) == (128, 6) - - -def _default_cfg(**overrides): - train_overrides = { - k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["train"] - } - model_overrides = { - k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["model"] - } - router_overrides = overrides.get("router") - if router_overrides: - model_overrides["router"] = router_overrides - return gconfig.merge_cli_overrides( - gconfig.DEFAULT_CONFIG, None, train_overrides, model_overrides - ) - - -_NOW = datetime(2026, 7, 29, 14, 30) - - -def test_default_out_dir_name_all_defaults_is_just_the_timestamp(): - cfg = _default_cfg() - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_single_non_default_field(): - cfg = _default_cfg(hidden_dim=512) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_h512" - - -def test_default_out_dir_name_conditioning_embedding_shown_abbreviated(): - cfg = _default_cfg(conditioning="embedding") - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb" - - -def test_default_out_dir_name_conditioning_default_omitted(): - cfg = _default_cfg(conditioning="physical") - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_router_enabled_shown_as_unit(): - cfg = _default_cfg(router={"enabled": True, "type": "energy", "n_experts": 8}) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8" - - -def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault(): - cfg = _default_cfg(router={"enabled": False, "type": "pdg", "n_experts": 8}) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_router_gumbel_shown_when_enabled(): - cfg = _default_cfg( - router={"enabled": True, "type": "energy", "n_experts": 8, "gumbel": True} - ) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_gum" - - -def test_default_out_dir_name_router_gumbel_omitted_when_router_disabled(): - cfg = _default_cfg( - router={"enabled": False, "type": "energy", "n_experts": 8, "gumbel": True} - ) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_router_learn_centers_shown_only_when_disabled(): - cfg_default = _default_cfg( - router={"enabled": True, "type": "energy", "n_experts": 8} - ) - assert ( - gconfig.default_out_dir_name(cfg_default, now=_NOW) == "20260729_1430_r-energy8" - ) - - cfg_off = _default_cfg( - router={ - "enabled": True, - "type": "energy", - "n_experts": 8, - "learn_centers": False, - } - ) - assert ( - gconfig.default_out_dir_name(cfg_off, now=_NOW) - == "20260729_1430_r-energy8_nolc" - ) - - -def test_default_out_dir_name_router_learn_width_and_temperature_shown(): - cfg = _default_cfg( - router={ - "enabled": True, - "type": "energy", - "n_experts": 8, - "learn_width": True, - } - ) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_lw" - - cfg2 = _default_cfg( - router={ - "enabled": True, - "type": "energy", - "n_experts": 8, - "learn_temperature": True, - } - ) - assert gconfig.default_out_dir_name(cfg2, now=_NOW) == "20260729_1430_r-energy8_lt" - - -def test_default_out_dir_name_mode_shown_bare_no_prefix(): - cfg = _default_cfg(mode="wgan") - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan" - - -def test_default_out_dir_name_overflow_caps_and_hashes_remainder(): - cfg = _default_cfg( - mode="wgan", - router={"enabled": True, "type": "energy", "n_experts": 8}, - conditioning="embedding", - hidden_dim=512, - n_blocks=8, - emb_dim=32, - lr=1e-3, - batch_size=2048, - seed=3, - epochs=200, - ) - name = gconfig.default_out_dir_name(cfg, now=_NOW) - # First 6 by priority: mode, router, conditioning, hidden_dim, n_blocks, emb_dim. - assert name.startswith("20260729_1430_wgan_r-energy8_cemb_h512_b8_e32_+4more-") - digest = name.split("-")[-1] - assert len(digest) == 6 - - -def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive(): - base = dict( - mode="wgan", - router={"enabled": True, "type": "energy", "n_experts": 8}, - conditioning="embedding", - hidden_dim=512, - n_blocks=8, - emb_dim=32, - lr=1e-3, - batch_size=2048, - seed=3, - epochs=200, - ) - name_a = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW) - name_b = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW) - assert name_a == name_b # stable across calls with the same overflow set - - changed = dict(base, epochs=999) - name_c = gconfig.default_out_dir_name(_default_cfg(**changed), now=_NOW) - assert ( - name_c != name_a - ) # differs when an overflowed value changes # n_blocks inherited, hidden_dim not