lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Make config dataclasses the single source of truth for DEFAULT_CONFIG
CI / Format (ruff format) (push) Successful in 31s
Details
CI / Lint (ruff check) (push) Successful in 31s
Details
CI / Sync project version with tag (push) Has been skipped
Details
CI / Lint (ruff check) (pull_request) Successful in 44s
Details
CI / Type check (ty) (push) Successful in 46s
Details
CI / Format (ruff format) (pull_request) Successful in 39s
Details
CI / Sync project version with tag (pull_request) Has been skipped
Details
CI / Type check (ty) (pull_request) Successful in 39s
Details
CI / Tests (pull_request) Successful in 3m38s
Details
CI / Tests (push) Successful in 3m45s
Details

Browse Source 
DEFAULT_CONFIG and build_models/build_critics/StageSpec.from_config's
inline .get(key, default) fallbacks had already drifted: two keys
(stage2_model.decoder, stage2_model.particle_type.target) resolved
differently depending on whether a config dict came from
merge_cli_overrides (fully populated, correct) or was hand-built and
partial (fell back to stale v0.2-shaped literals). Introduce frozen
dataclasses (GiantConfig and its nested blocks) in giant/config.py as
the actual single declaration of every default; DEFAULT_CONFIG is now
generated from them instead of hand-maintained, and build_models,
build_critics, and StageSpec.from_config consume the dataclasses
instead of duplicating literal fallbacks, so this class of drift can't
recur. Router/n_sec sub-blocks keep an `extra` catch-all for their
genuinely dynamic keys (composed-router axes, runtime-seeded
centers_init, legacy_owner).

Fixing the fallback surfaced the same latent bug in two existing
partial-config callers that had been silently depending on it: a
test fixture in test_train.py and scripts/warm_setup_cache.py's
minimal cfg (now merged against DEFAULT_CONFIG instead of hand-rolled,
closing the gap for good). See issues.md Issue 1.

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