v0.3.0 step 1: new nested config schema, v0.2 migration shim
Replace the single global train.mode + [model] block with the four top-level blocks docs/v0.3.0-design.md specifies ([conditioning], [stage1_model], [stage2_model], [train]), so Stage 1 and Stage 2 can run independent generative objectives and Stage 2 can train standalone. - migrate_config translates old config.toml/checkpoint dicts on load, so nothing on /ceph goes dead; loudly rejects non-zero expert_hidden_dim/expert_n_blocks, which v0.3.0 no longer supports. - merge_cli_overrides/save_config generalize from one hardcoded nesting level (model.router) to arbitrary recursive depth. - default_out_dir_name candidates move to dotted paths against the new schema, with per-stage router/generator discriminators. - validate_config adds cross-block checks the per-block schema can't express (particle_type.target=embedding needs a matching conditioning mode, tie_to_stage1 needs an active stage 1, etc). - resolve_expert_dims is deleted (experts always inherit the stage's hidden_dim/n_res_blocks now) — pipeline.py/cli.py callers are left dangling on purpose, to be updated in the network.py/train.py steps that follow. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
+661
-217
@@ -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"),
|
||||
|
||||
+575
-222
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user