lars
87e37ebe14
CI / Format (ruff format) (push) Successful in 29s
CI / Lint (ruff check) (push) Successful in 30s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Lint (ruff check) (pull_request) Successful in 31s
CI / Format (ruff format) (pull_request) Successful in 44s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 4m32s
CI / Tests (pull_request) Successful in 4m25s
stage{1,2}_model.active = false already trains one stage alone, but the
checkpoint it writes holds only that stage, so giant rollout refuses it --
the "retrain stage 2 alone against a fixed, known-good stage 1" experiment
the 2026-08-03 species failure calls for wasn't runnable end to end.
Adds stage{1,2}_model.init_from (a checkpoint .pt to load this stage's
weights from before training) and .freeze (never update them), symmetric
across both stages. Both stages stay active = true, so both get built and
both land in the output checkpoint -- the frozen stage is merely
initialized from disk instead of from scratch.
Decisions made during planning:
- Soft freeze: forward/backward still run every batch (loss/grad_norm stay
meaningful, no autograd special-casing), only optimizer.step() (and, for
the frozen stage, lr_sched.step()/EMA update) is skipped -- weights are
byte-identical for the whole run. This is StageTrainer._step_optimizer,
shared by the non-adversarial path and both halves (generator + critic)
of the WGAN path, so a frozen WGAN stage's critic freezes too.
- validate_config requires init_from whenever freeze = true, unless the run
is a --resume (a resumed frozen stage's weights come from the resume
checkpoint instead) -- freezing a randomly-initialized model is almost
certainly a mistake.
- CLI flags on both `giant train` and `giant new-run`
(--stage{1,2}-init-from/--stage{1,2}-freeze), matching every other
per-stage model knob's existing treatment.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
1791 lines
75 KiB
Python
1791 lines
75 KiB
Python
import copy
|
|
import difflib
|
|
import hashlib
|
|
import random
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tomllib
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
|
|
from giant.model.history import HISTORY_REGISTRY
|
|
|
|
|
|
class Conditioning(str, Enum):
|
|
"""`conditioning.particle.type` / `conditioning.material.type` choices —
|
|
shared by `giant.cli` and `giant.tools.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
|
|
|
|
|
|
# --- 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`.
|
|
#
|
|
# The router sub-block carries 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`. It carries 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 TrunkConfig:
|
|
"""`stage1_model.trunk`/`stage2_model.trunk`: selects the trunk's expert
|
|
*body* architecture from `giant.model.trunks.TRUNK_REGISTRY` (default
|
|
`"resmlp"` — today's only body, `input_proj -> ResBlock stack ->
|
|
out_proj`). Orthogonal to whether that body is mixed: mixing is still
|
|
controlled entirely by `router.enabled`/`router.n_experts` on the same
|
|
stage, unaffected by this block. A future body's own hyperparameters
|
|
(e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling
|
|
field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by
|
|
`generator`.
|
|
|
|
`block_conditioning` selects each body's conditioning-injection mechanism
|
|
from `giant.model.layers.BLOCK_REGISTRY` — `"add"` (default, today's
|
|
conditional-bias `ResBlock`, bit-identical to pre-gitea-#34 behaviour),
|
|
`"film"`, or `"adaln"`."""
|
|
|
|
type: str = "resmlp"
|
|
block_conditioning: str = "add"
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict | None) -> "TrunkConfig":
|
|
d = d or {}
|
|
return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))
|
|
|
|
def to_dict(self) -> dict:
|
|
return {"type": self.type, "block_conditioning": self.block_conditioning}
|
|
|
|
|
|
@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 per-slot stop head on the autoregressive
|
|
# secondary decoder (Stage2Autoregressive only — see validate_config),
|
|
# evaluated against the generated prefix instead of conditioning alone.
|
|
# Replaces n_sec_head entirely: the two are mutually exclusive, so this
|
|
# mode builds no n_sec_head and stage2_model.n_sec.lambda instead weights
|
|
# the stop head's BCE term.
|
|
# "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/BCE weight for the head
|
|
# Which stage's module physically owns the n_sec_head weights: "stage2" (default,
|
|
# fresh v0.3.0 runs — Stage2OneShot/Stage2Autoregressive builds it) or "stage1"
|
|
# (a migrated v0.2 checkpoint — see network._migrate_legacy_model_config, whose
|
|
# n_sec head was trained against Stage 1's own ConditionEncoder output and so has
|
|
# to stay attached there, not just be labeled as such).
|
|
owner: str = "stage2"
|
|
# mode="stop_token" only: how sample_secondaries_ar turns a slot's stop logit into a
|
|
# stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample":
|
|
# a Bernoulli draw at sigmoid(logit) (a real sample from the learned length
|
|
# distribution, at the cost of an extra RNG draw per slot).
|
|
stop_sampling: str = "greedy"
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict | None) -> "NSecConfig":
|
|
d = d or {}
|
|
return cls(
|
|
mode=d.get("mode", "head"),
|
|
lambda_weight=d.get("lambda", 0.1),
|
|
owner=d.get("owner", "stage2"),
|
|
stop_sampling=d.get("stop_sampling", "greedy"),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"mode": self.mode,
|
|
"lambda": self.lambda_weight,
|
|
"owner": self.owner,
|
|
"stop_sampling": self.stop_sampling,
|
|
}
|
|
|
|
|
|
@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"
|
|
# Secondary-species class count under target = "onehot" — independent of
|
|
# conditioning.particle.emb_dim (see gitea #29: the two used to be
|
|
# silently the same number). 0 = inherit conditioning.particle.emb_dim,
|
|
# preserving pre-#29 behavior.
|
|
n_classes: int = 0
|
|
|
|
@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"),
|
|
n_classes=d.get("n_classes", 0),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"target": self.target,
|
|
"lambda": self.lambda_weight,
|
|
"other_policy": self.other_policy,
|
|
"n_classes": self.n_classes,
|
|
}
|
|
|
|
|
|
@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 HeadConfig:
|
|
"""A single classifier head's shape — `n_sec_head`/`type_head` (gitea
|
|
#36 deduplicated their five identical hand-rolled
|
|
`Linear -> SiLU -> Linear` definitions into
|
|
`giant.model.layers.build_mlp_head`, which this config drives).
|
|
`hidden_ratio=0.5`/`depth=2` are the exact pre-#36 hardcoded values
|
|
(hidden width = `hidden_dim // 2`, one hidden layer), so omitting a
|
|
`heads` block — including every migrated v0.2 config — reproduces the
|
|
old architecture bit-for-bit."""
|
|
|
|
hidden_ratio: float = 0.5 # hidden width = round(hidden_dim * hidden_ratio)
|
|
depth: int = 2 # matches build_mlp_head's depth
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict | None) -> "HeadConfig":
|
|
d = d or {}
|
|
return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2))
|
|
|
|
def to_dict(self) -> dict:
|
|
return {"hidden_ratio": self.hidden_ratio, "depth": self.depth}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Stage1HeadsConfig:
|
|
"""Stage 1 only ever owns `n_sec_head`, and only for a migrated v0.2
|
|
checkpoint (`stage2_model.n_sec.owner = "stage1"`) — see
|
|
`Stage1Model`'s docstring."""
|
|
|
|
n_sec: HeadConfig = field(default_factory=HeadConfig)
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig":
|
|
d = d or {}
|
|
return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")))
|
|
|
|
def to_dict(self) -> dict:
|
|
return {"n_sec": self.n_sec.to_dict()}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Stage2HeadsConfig:
|
|
"""`n_sec` and `type` are independently configurable — n_sec accuracy
|
|
and secondary-species accuracy are separately known weak spots (gitea
|
|
#36)."""
|
|
|
|
n_sec: HeadConfig = field(default_factory=HeadConfig)
|
|
type: HeadConfig = field(default_factory=HeadConfig)
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig":
|
|
d = d or {}
|
|
return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type")))
|
|
|
|
def to_dict(self) -> dict:
|
|
return {"n_sec": self.n_sec.to_dict(), "type": self.type.to_dict()}
|
|
|
|
|
|
@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
|
|
# Checkpoint .pt to load this stage's weights from before training starts
|
|
# (its own "model"/"sec_decoder" key, not this run's own resume state) —
|
|
# "" means start from a fresh init. See `freeze` below for the partial-
|
|
# retrain use case this exists for (gitea #42).
|
|
init_from: str = ""
|
|
# true keeps this stage's weights exactly as loaded from `init_from` —
|
|
# forward/backward still run every batch (so its loss/grad_norm metrics
|
|
# stay meaningful, and a WGAN stage's critic still gets a real signal to
|
|
# report), but its optimizer never steps. Lets a rollout-capable
|
|
# checkpoint retrain only the *other* stage against a fixed, known-good
|
|
# one (gitea #42) — `validate_config` requires `init_from` to be set
|
|
# whenever this is true, unless the run is a `--resume`.
|
|
freeze: bool = False
|
|
# "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)
|
|
trunk: TrunkConfig = field(default_factory=TrunkConfig)
|
|
heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig)
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict | None) -> "Stage1ModelConfig":
|
|
d = d or {}
|
|
return cls(
|
|
active=d.get("active", True),
|
|
init_from=d.get("init_from", ""),
|
|
freeze=d.get("freeze", False),
|
|
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")),
|
|
trunk=TrunkConfig.from_dict(d.get("trunk")),
|
|
heads=Stage1HeadsConfig.from_dict(d.get("heads")),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"active": self.active,
|
|
"init_from": self.init_from,
|
|
"freeze": self.freeze,
|
|
"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(),
|
|
"trunk": self.trunk.to_dict(),
|
|
"heads": self.heads.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
|
|
# See Stage1ModelConfig.init_from/.freeze — same semantics, this stage.
|
|
init_from: str = ""
|
|
freeze: bool = False
|
|
# "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"
|
|
# Ramp for "sampled": P(condition on the ground-truth stage-1 outcome
|
|
# rather than a fresh sample), linearly interpolated from ctx_p_start
|
|
# (epoch 0) to ctx_p_end (the final epoch) — the same scheduled-sampling
|
|
# shape as autoregressive.tf_p_start/tf_p_end, so stage 2 doesn't chase a
|
|
# wildly moving stage-1 target in early epochs. Unread under "truth".
|
|
ctx_p_start: float = 1.0
|
|
ctx_p_end: float = 0.0
|
|
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)
|
|
trunk: TrunkConfig = field(default_factory=TrunkConfig)
|
|
heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig)
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict | None) -> "Stage2ModelConfig":
|
|
d = d or {}
|
|
return cls(
|
|
active=d.get("active", True),
|
|
init_from=d.get("init_from", ""),
|
|
freeze=d.get("freeze", False),
|
|
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"),
|
|
ctx_p_start=d.get("ctx_p_start", 1.0),
|
|
ctx_p_end=d.get("ctx_p_end", 0.0),
|
|
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")),
|
|
trunk=TrunkConfig.from_dict(d.get("trunk")),
|
|
heads=Stage2HeadsConfig.from_dict(d.get("heads")),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"active": self.active,
|
|
"init_from": self.init_from,
|
|
"freeze": self.freeze,
|
|
"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,
|
|
"ctx_p_start": self.ctx_p_start,
|
|
"ctx_p_end": self.ctx_p_end,
|
|
"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(),
|
|
"trunk": self.trunk.to_dict(),
|
|
"heads": self.heads.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 leaf_paths(node: dict, prefix: str = "") -> list[str]:
|
|
"""Every dotted leaf path in a DEFAULT_CONFIG-shaped dict, e.g.
|
|
"stage1_model.router.n_experts". `[meta]` (run provenance, no schema
|
|
counterpart) is skipped at the top level, matching `validate_config_keys`.
|
|
Shared by `tests/test_config_consumed_keys.py` (the static per-identifier
|
|
audit) and `giant.model.summary` (the runtime per-config audit, gitea
|
|
#46) so both walk the exact same tree."""
|
|
paths = []
|
|
for key, value in node.items():
|
|
if prefix == "" and key == "meta":
|
|
continue
|
|
path = f"{prefix}.{key}" if prefix else key
|
|
if isinstance(value, dict):
|
|
paths.extend(leaf_paths(value, path))
|
|
else:
|
|
paths.append(path)
|
|
return paths
|
|
|
|
|
|
def git_hash() -> str:
|
|
try:
|
|
return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
|
|
except Exception:
|
|
return "unknown"
|
|
|
|
|
|
def auto_device() -> torch.device:
|
|
if torch.cuda.is_available():
|
|
return torch.device("cuda")
|
|
if torch.backends.mps.is_available():
|
|
return torch.device("mps")
|
|
return torch.device("cpu")
|
|
|
|
|
|
# Calibration point for estimate_batch_size(training=True): hidden_dim=1024,
|
|
# n_blocks=6, batch_size=29696 measured at ~7683 MiB VRAM (post-Phase-2
|
|
# architecture, including the Stage-2 secondary decoder and n_sec head).
|
|
# 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 —
|
|
# deliberately last in the implementation order.
|
|
_REF_BYTES = 7683 * 1024**2
|
|
_REF_BATCH_SIZE = 29696
|
|
_REF_HIDDEN_DIM = 1024
|
|
_REF_N_BLOCKS = 6
|
|
|
|
# Calibration point for estimate_batch_size(training=False): inference has no
|
|
# backward graph or optimizer state, so its memory footprint is much smaller
|
|
# per sample. hidden_dim=1024, n_blocks=8, batch_size=65536 measured at ~2037
|
|
# MiB VRAM.
|
|
_REF_BYTES_PREDICT = 2037 * 1024**2
|
|
_REF_BATCH_SIZE_PREDICT = 65536
|
|
_REF_HIDDEN_DIM_PREDICT = 1024
|
|
_REF_N_BLOCKS_PREDICT = 8
|
|
|
|
|
|
def estimate_batch_size(
|
|
hidden_dim: int,
|
|
n_blocks: int,
|
|
device: torch.device,
|
|
safety_factor: float = 0.8,
|
|
min_batch_size: int = 1024,
|
|
training: bool = True,
|
|
) -> int:
|
|
"""Estimate a batch size that fits in the free memory on `device`.
|
|
|
|
Only supported on CUDA devices, which expose a free/total memory query;
|
|
other backends (cpu, mps) raise ValueError. Pass `training=False` for
|
|
inference (e.g. `predict`), which uses a much lower per-sample memory
|
|
calibration since there's no backward graph or optimizer state.
|
|
"""
|
|
if device.type != "cuda":
|
|
raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}")
|
|
device_index = device.index if device.index is not None else torch.cuda.current_device()
|
|
free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
|
|
if training:
|
|
ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
|
|
_REF_BYTES,
|
|
_REF_BATCH_SIZE,
|
|
_REF_HIDDEN_DIM,
|
|
_REF_N_BLOCKS,
|
|
)
|
|
else:
|
|
ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
|
|
_REF_BYTES_PREDICT,
|
|
_REF_BATCH_SIZE_PREDICT,
|
|
_REF_HIDDEN_DIM_PREDICT,
|
|
_REF_N_BLOCKS_PREDICT,
|
|
)
|
|
bytes_per_unit = ref_bytes / (ref_batch_size * ref_hidden_dim * ref_n_blocks)
|
|
bytes_per_sample = bytes_per_unit * hidden_dim * n_blocks
|
|
batch_size = int(free_bytes * safety_factor / bytes_per_sample)
|
|
batch_size = max(min_batch_size, (batch_size // 1024) * 1024)
|
|
return batch_size
|
|
|
|
|
|
def load_toml(path: Path) -> dict:
|
|
with open(path, "rb") as f:
|
|
return tomllib.load(f)
|
|
|
|
|
|
def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None:
|
|
"""Warn (don't fail) if a config.toml's [meta].git_hash predates the current checkout.
|
|
|
|
A config saved by a previous run may have been produced by code that has
|
|
since changed, so its hyperparameters might not mean what they used to —
|
|
surface that as a heads-up rather than blocking the rerun.
|
|
"""
|
|
file_hash = file_cfg.get("meta", {}).get("git_hash")
|
|
current_hash = git_hash()
|
|
if not file_hash or file_hash == "unknown" or current_hash == "unknown":
|
|
return
|
|
if file_hash != current_hash:
|
|
print(
|
|
f"warning: {config_path} was generated at git commit {file_hash}, "
|
|
f"but the current checkout is at {current_hash} — hyperparameters "
|
|
"may not match the code that originally produced this config",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
|
|
def load_checkpoint_config(ckpt_path: str | Path) -> dict:
|
|
"""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. 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():
|
|
return {}
|
|
return load_toml(config_path)
|
|
|
|
|
|
def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None:
|
|
"""Look for a config.toml next to a checkpoint and warn on a git_hash mismatch.
|
|
|
|
Training writes config.toml into the same out_dir as its checkpoints, so a
|
|
checkpoint loaded later (for `predict` or `giant.analysis`) can be
|
|
cross-checked the same way `--config` loading is, without the caller having
|
|
to pass the toml path explicitly. Silently does nothing if no config.toml
|
|
is found alongside the checkpoint.
|
|
"""
|
|
config_path = Path(ckpt_path).parent / "config.toml"
|
|
if not config_path.exists():
|
|
return
|
|
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
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FlagSpec:
|
|
"""One CLI flag's mapping into the config-overrides tree.
|
|
|
|
`paths` lists every dotted config path this flag writes (>1 means fan-out
|
|
to multiple stages/axes, e.g. `--mode` -> both stages' `generator`).
|
|
`precedence` controls write order when two flags target the same path:
|
|
specs are applied in ascending precedence, so a higher-precedence (more
|
|
specific) flag overwrites a lower-precedence (shared/shorthand) one —
|
|
this is the "build a shared dict, then let a more specific dict win"
|
|
pattern `giant train`/`giant new-run` need (e.g. `--hidden-dim` vs
|
|
`--stage1-hidden-dim`, or `--n-critic` vs `--stage1-n-critic`),
|
|
generalized to one mechanism instead of three different ad hoc ones.
|
|
"""
|
|
|
|
name: str
|
|
paths: tuple[str, ...]
|
|
precedence: int = 0
|
|
|
|
|
|
# Flag -> config-path table shared by `giant train`/`giant new-run`
|
|
# (giant/cli.py) so both commands resolve CLI overrides identically. See
|
|
# issues.md Issue 3: this replaces ~140 lines of hand-written, imperative
|
|
# dict-building in cli.py with one declarative table plus
|
|
# `overrides_from_flags` below.
|
|
FLAG_SPECS: tuple[FlagSpec, ...] = (
|
|
# train block -- flat pass-through, unique paths, precedence irrelevant.
|
|
FlagSpec("epochs", ("train.epochs",)),
|
|
FlagSpec("batch_size", ("train.batch_size",)),
|
|
FlagSpec("lr", ("train.lr",)),
|
|
FlagSpec("weight_decay", ("train.weight_decay",)),
|
|
FlagSpec("ema_decay", ("train.ema_decay",)),
|
|
FlagSpec("warmup_epochs", ("train.warmup_epochs",)),
|
|
FlagSpec("val_fraction", ("train.val_fraction",)),
|
|
FlagSpec("num_workers", ("train.num_workers",)),
|
|
FlagSpec("seed", ("train.seed",)),
|
|
FlagSpec("validate_every", ("train.validate_every",)),
|
|
FlagSpec("validate_steps", ("train.validate_steps",)),
|
|
FlagSpec("max_val_batches", ("train.max_val_batches",)),
|
|
FlagSpec("wandb", ("train.wandb",)),
|
|
FlagSpec("wandb_project", ("train.wandb_project",)),
|
|
FlagSpec("wandb_run_name", ("train.wandb_run_name",)),
|
|
FlagSpec("wandb_log_every", ("train.wandb_log_every",)),
|
|
# --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat
|
|
# shorthands (they predate stage2_model having its own flags);
|
|
# --stage1-* wins when both are given.
|
|
FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0),
|
|
FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1),
|
|
FlagSpec("n_blocks", ("stage1_model.n_res_blocks",), precedence=0),
|
|
FlagSpec("stage1_n_res_blocks", ("stage1_model.n_res_blocks",), precedence=1),
|
|
FlagSpec("dropout", ("stage1_model.dropout",), precedence=0),
|
|
FlagSpec("stage1_dropout", ("stage1_model.dropout",), precedence=1),
|
|
# stage2-only knobs.
|
|
FlagSpec("stage2_hidden_dim", ("stage2_model.hidden_dim",)),
|
|
FlagSpec("stage2_n_res_blocks", ("stage2_model.n_res_blocks",)),
|
|
FlagSpec("stage2_dropout", ("stage2_model.dropout",)),
|
|
FlagSpec("stage2_decoder", ("stage2_model.decoder",)),
|
|
FlagSpec("stage2_k_max", ("stage2_model.k_max",)),
|
|
FlagSpec("stage2_context_dim", ("stage2_model.context_dim",)),
|
|
FlagSpec("stage2_stage1_context", ("stage2_model.stage1_context",)),
|
|
# --mode applies to both stages by default (v0.2 had one shared
|
|
# mode/wgan config); --stage{1,2}-generator override a single stage.
|
|
FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator"), precedence=0),
|
|
FlagSpec("stage1_generator", ("stage1_model.generator",), precedence=1),
|
|
FlagSpec("stage2_generator", ("stage2_model.generator",), precedence=1),
|
|
# --emb-dim/--conditioning set both conditioning axes (v0.2 had one
|
|
# shared value for particle+material).
|
|
FlagSpec("conditioning", ("conditioning.particle.type", "conditioning.material.type")),
|
|
FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")),
|
|
# Pre-aggregated router override dict (built by `_router_cli_overrides`
|
|
# in cli.py from --router/--router-type/--n-experts/--router-axis).
|
|
# Router overrides only ever land on stage1_model -- this asymmetry is
|
|
# deliberate (see cli.py) and must not be "fixed" into a fan-out here.
|
|
FlagSpec("router_config", ("stage1_model.router",)),
|
|
# WGAN: shared knobs apply to both stages by default (v0.2 had one
|
|
# shared wgan config); --stage{1,2}-* override a single stage.
|
|
FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic"), precedence=0),
|
|
FlagSpec("stage1_n_critic", ("stage1_model.wgan.n_critic",), precedence=1),
|
|
FlagSpec("stage2_n_critic", ("stage2_model.wgan.n_critic",), precedence=1),
|
|
FlagSpec("gp_weight", ("stage1_model.wgan.gp_weight", "stage2_model.wgan.gp_weight"), precedence=0),
|
|
FlagSpec("stage1_gp_weight", ("stage1_model.wgan.gp_weight",), precedence=1),
|
|
FlagSpec("stage2_gp_weight", ("stage2_model.wgan.gp_weight",), precedence=1),
|
|
FlagSpec("noise_dim", ("stage1_model.wgan.noise_dim", "stage2_model.wgan.noise_dim"), precedence=0),
|
|
FlagSpec("stage1_noise_dim", ("stage1_model.wgan.noise_dim",), precedence=1),
|
|
FlagSpec("stage2_noise_dim", ("stage2_model.wgan.noise_dim",), precedence=1),
|
|
FlagSpec("critic_lr", ("stage1_model.wgan.critic_lr", "stage2_model.wgan.critic_lr"), precedence=0),
|
|
FlagSpec("stage1_critic_lr", ("stage1_model.wgan.critic_lr",), precedence=1),
|
|
FlagSpec("stage2_critic_lr", ("stage2_model.wgan.critic_lr",), precedence=1),
|
|
# Critic sizing: stage-scoped only, no shared alias — this is an
|
|
# architectural per-stage knob like hidden_dim/n_res_blocks above, not a
|
|
# shared training hyperparameter like the wgan knobs above it.
|
|
FlagSpec("stage1_critic_hidden_dim", ("stage1_model.wgan.critic_hidden_dim",)),
|
|
FlagSpec("stage1_critic_n_res_blocks", ("stage1_model.wgan.critic_n_res_blocks",)),
|
|
FlagSpec("stage2_critic_hidden_dim", ("stage2_model.wgan.critic_hidden_dim",)),
|
|
FlagSpec("stage2_critic_n_res_blocks", ("stage2_model.wgan.critic_n_res_blocks",)),
|
|
# Partial-retrain (gitea #42): stage-scoped only, no shared alias — a
|
|
# shared "freeze both stages from the same file" flag has no sensible
|
|
# meaning (a checkpoint has one set of weights per stage).
|
|
FlagSpec("stage1_init_from", ("stage1_model.init_from",)),
|
|
FlagSpec("stage1_freeze", ("stage1_model.freeze",)),
|
|
FlagSpec("stage2_init_from", ("stage2_model.init_from",)),
|
|
FlagSpec("stage2_freeze", ("stage2_model.freeze",)),
|
|
)
|
|
|
|
|
|
def overrides_from_flags(values: dict[str, object]) -> dict:
|
|
"""Build the nested, section-keyed config-overrides dict
|
|
`merge_cli_overrides` expects, from `{flag_name: value}`.
|
|
|
|
Flags absent from `values`, or mapped to `None` (= not given on the
|
|
CLI), are skipped. See `FlagSpec`/`FLAG_SPECS` above for the precedence
|
|
rule applied when two flags target the same path.
|
|
"""
|
|
overrides: dict = {}
|
|
for spec in sorted(FLAG_SPECS, key=lambda s: s.precedence):
|
|
if spec.name not in values or values[spec.name] is None:
|
|
continue
|
|
for path in spec.paths:
|
|
_set_path(overrides, path, values[spec.name])
|
|
return overrides
|
|
|
|
|
|
# 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 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. `[meta].config_version == CONFIG_VERSION` marks a dict as
|
|
already-v0.3; its absence is read as "this is v0.2", 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 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_KEY_TO_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:
|
|
reject_legacy_router_expert_sizing(old_router, source="v0.2 config's model.router")
|
|
_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 (see giant._migration).
|
|
for path, value in V02_FIXED_FACTS.items():
|
|
_set_path(new, path, value)
|
|
|
|
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
|
|
|
|
|
|
# axis{i}_{field} composed-router keys (see network._parse_composed_axes) — field-name
|
|
# agnostic, matching network.py's own _AXIS_KEY_RE, since anything after axis{i}_ is
|
|
# passed straight through as a router kwarg there.
|
|
_AXIS_KEY_RE = re.compile(r"^axis\d+_.+$")
|
|
|
|
# Paths (dotted, relative to the merged cfg root) that carry genuinely dynamic keys not
|
|
# in DEFAULT_CONFIG's fixed schema — composed-router axis{i}_{field} flags and
|
|
# pipeline.seed_router_centers's runtime-seeded centers_init (see RouterConfig.extra
|
|
# above). validate_config_keys allows any key under these paths through unconditionally
|
|
# apart from the axis-pattern/centers_init check below.
|
|
_DYNAMIC_ROUTER_PATHS = {"stage1_model.router", "stage2_model.router"}
|
|
|
|
|
|
def _unknown_key_error(full_path: str, key: str, valid_keys) -> ValueError:
|
|
hint = difflib.get_close_matches(key, list(valid_keys), n=1)
|
|
suggestion = f" — did you mean {hint[0]!r}?" if hint else ""
|
|
return ValueError(f"unknown config key {full_path!r}{suggestion}")
|
|
|
|
|
|
def _validate_keys(node: dict, default_node: dict, path: str) -> None:
|
|
for key, value in node.items():
|
|
if path == "" and key == "meta":
|
|
continue
|
|
full_path = f"{path}.{key}" if path else key
|
|
if path in _DYNAMIC_ROUTER_PATHS and (key == "centers_init" or _AXIS_KEY_RE.match(key)):
|
|
continue
|
|
if key not in default_node:
|
|
raise _unknown_key_error(full_path, key, default_node.keys())
|
|
if isinstance(value, dict) and isinstance(default_node[key], dict):
|
|
_validate_keys(value, default_node[key], full_path)
|
|
|
|
|
|
def validate_config_keys(cfg: dict) -> None:
|
|
"""Reject any config key not part of the known v0.3 schema (DEFAULT_CONFIG's tree).
|
|
|
|
Catches typos like `n_res_block` for `n_res_blocks` that would otherwise merge
|
|
cleanly, pass `validate_config`, and silently build the wrong model — see
|
|
issues.md Issue 2.
|
|
|
|
Only exercised on the config.toml/CLI-overrides path (called from
|
|
`merge_cli_overrides` below). A checkpoint's `model_config` dict goes through
|
|
`network._migrate_legacy_model_config`/`build_models` instead and must keep
|
|
loading regardless of schema drift; old checkpoints predate this validator and
|
|
are never passed through here.
|
|
|
|
Two areas are deliberately dynamic and excluded: `[meta]` (run provenance, no
|
|
DEFAULT_CONFIG counterpart), and `stage{1,2}_model.router`'s `axis{i}_{field}`
|
|
keys (composed-router axes, see `network._parse_composed_axes`) /
|
|
`centers_init` (runtime-seeded by `pipeline.seed_router_centers`).
|
|
"""
|
|
_validate_keys(cfg, DEFAULT_CONFIG, "")
|
|
|
|
|
|
def merge_cli_overrides(
|
|
defaults: dict,
|
|
config_path: Path | None,
|
|
overrides: dict,
|
|
) -> dict:
|
|
"""Resolve config as defaults -> TOML file -> explicit overrides.
|
|
|
|
`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. The result is checked against the known schema
|
|
(`validate_config_keys`) before being returned, so a typo'd key raises
|
|
here rather than silently building the wrong model.
|
|
"""
|
|
cfg = copy.deepcopy(defaults)
|
|
if config_path is not None:
|
|
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)
|
|
for section, values in overrides.items():
|
|
if isinstance(values, dict):
|
|
cfg[section] = _deep_merge(cfg.get(section, {}), values)
|
|
else:
|
|
cfg[section] = values
|
|
validate_config_keys(cfg)
|
|
return cfg
|
|
|
|
|
|
def validate_config(cfg: dict, *, resume: bool = False) -> None:
|
|
"""Cross-block validation the per-block schema can't express on its own.
|
|
|
|
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.
|
|
|
|
`resume=True` (only `giant train --resume` passes this) relaxes the
|
|
`stage{1,2}_model.freeze` -> `.init_from` requirement below: a resumed
|
|
frozen stage's weights come from the resume checkpoint, not `init_from`.
|
|
"""
|
|
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"):
|
|
if _get_path(cfg, f"{stage_name}.freeze") and not _get_path(cfg, f"{stage_name}.init_from") and not resume:
|
|
raise ValueError(
|
|
f"{stage_name}.freeze = true requires {stage_name}.init_from "
|
|
"to be set (or --resume) — freezing a randomly-initialized "
|
|
"model is almost certainly a mistake"
|
|
)
|
|
|
|
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":
|
|
if _get_path(cfg, "stage2_model.decoder") != "autoregressive":
|
|
raise ValueError(
|
|
"stage2_model.n_sec.mode = 'stop_token' requires "
|
|
"stage2_model.decoder = 'autoregressive' — there is no "
|
|
"per-token loop to stop under 'one_shot'"
|
|
)
|
|
if _get_path(cfg, "stage2_model.n_sec.owner") != "stage2":
|
|
raise ValueError(
|
|
"stage2_model.n_sec.mode = 'stop_token' requires "
|
|
"stage2_model.n_sec.owner = 'stage2' — a migrated v0.2 "
|
|
"checkpoint's stage-1 n_sec_head has no per-token "
|
|
"conditioning to hang an EOS decision off"
|
|
)
|
|
|
|
stop_sampling = _get_path(cfg, "stage2_model.n_sec.stop_sampling")
|
|
if stop_sampling not in ("greedy", "sample"):
|
|
raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
|
|
|
|
stage1_context = _get_path(cfg, "stage2_model.stage1_context")
|
|
if stage1_context not in ("truth", "sampled"):
|
|
raise ValueError(f"stage2_model.stage1_context = {stage1_context!r} — must be 'truth' or 'sampled'")
|
|
if stage1_context == "sampled":
|
|
if not (_get_path(cfg, "stage1_model.active") and _get_path(cfg, "stage2_model.active")):
|
|
raise ValueError(
|
|
"stage2_model.stage1_context = 'sampled' requires both "
|
|
"stage1_model.active and stage2_model.active = true — there is "
|
|
"no stage-1 model to sample from in a stage-2-only run"
|
|
)
|
|
ctx_p_start = _get_path(cfg, "stage2_model.ctx_p_start")
|
|
ctx_p_end = _get_path(cfg, "stage2_model.ctx_p_end")
|
|
for name, value in (("ctx_p_start", ctx_p_start), ("ctx_p_end", ctx_p_end)):
|
|
if not (0.0 <= value <= 1.0):
|
|
raise ValueError(f"stage2_model.{name} = {value} — must be in [0, 1]")
|
|
if ctx_p_start == 1.0 and ctx_p_end == 1.0:
|
|
raise ValueError(
|
|
"stage2_model.stage1_context = 'sampled' with ctx_p_start = "
|
|
"ctx_p_end = 1.0 always conditions on the ground truth — "
|
|
"identical to 'truth' but silently so; use 'truth' instead or "
|
|
"lower ctx_p_end"
|
|
)
|
|
|
|
if (
|
|
_get_path(cfg, "stage2_model.n_sec.mode") == "truth"
|
|
and _get_path(cfg, "stage1_model.active")
|
|
and _get_path(cfg, "stage2_model.active")
|
|
):
|
|
raise ValueError(
|
|
"stage2_model.n_sec.mode = 'truth' is invalid for a "
|
|
"rollout-capable checkpoint (both stage1_model.active and "
|
|
"stage2_model.active = true): "
|
|
"'truth' takes n_sec from ground truth, which giant rollout "
|
|
"doesn't have. 'truth' is for standalone stage-2 evaluation "
|
|
"only — set stage1_model.active = false for that, or use "
|
|
"'head' (default) for a rollout-capable checkpoint."
|
|
)
|
|
|
|
if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
|
|
order = _get_path(cfg, "stage2_model.autoregressive.order")
|
|
if order != "energy_desc":
|
|
raise ValueError(
|
|
f"stage2_model.autoregressive.order = {order!r} — must be "
|
|
"'energy_desc' (the only implemented ordering; see "
|
|
"AutoregressiveConfig.order's docstring)"
|
|
)
|
|
history = _get_path(cfg, "stage2_model.autoregressive.history")
|
|
if history not in HISTORY_REGISTRY:
|
|
raise ValueError(
|
|
f"stage2_model.autoregressive.history = {history!r} — must be one of {sorted(HISTORY_REGISTRY)}"
|
|
)
|
|
teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
|
|
if teacher_forcing not in ("always", "scheduled", "never"):
|
|
raise ValueError(
|
|
"stage2_model.autoregressive.teacher_forcing = "
|
|
f"{teacher_forcing!r} — must be 'always', 'scheduled' or "
|
|
"'never'"
|
|
)
|
|
|
|
|
|
_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb", "onehot": "oh"}
|
|
|
|
|
|
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 _candidate(cfg):
|
|
value = _get_path(cfg, dotted_path)
|
|
default = _get_path(DEFAULT_CONFIG, dotted_path)
|
|
if value == default:
|
|
return None
|
|
return f"{prefix}{formatter(value)}"
|
|
|
|
return _candidate
|
|
|
|
|
|
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
|
|
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 = [
|
|
("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_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")),
|
|
("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")),
|
|
("stage1_block_cond", _path_candidate("stage1_model.trunk.block_conditioning", "s1bc-")),
|
|
("stage2_block_cond", _path_candidate("stage2_model.trunk.block_conditioning", "s2bc-")),
|
|
("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_freeze", _path_candidate("stage1_model.freeze", "s1frozen", formatter=lambda _: "")),
|
|
("stage2_freeze", _path_candidate("stage2_model.freeze", "s2frozen", formatter=lambda _: "")),
|
|
("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
|
|
|
|
|
|
def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
|
|
"""Build a default checkpoint out_dir name from what's non-default in `cfg`.
|
|
|
|
Only fields that differ from DEFAULT_CONFIG are included, so a fully
|
|
default run's name is just its timestamp — see
|
|
`_OUT_DIR_NAME_CANDIDATES` for the fixed, priority-ordered field list.
|
|
Beyond `_OUT_DIR_NAME_MAX_FIELDS` non-default fields, the remainder
|
|
collapse into a short deterministic hash suffix rather than growing the
|
|
name unboundedly. This name doubles as the run's W&B id (see
|
|
giant.training), which is the reason a timestamp is always included.
|
|
"""
|
|
now = now or datetime.now()
|
|
tokens = []
|
|
overflow = []
|
|
for label, candidate in _OUT_DIR_NAME_CANDIDATES:
|
|
token = candidate(cfg)
|
|
if token is None:
|
|
continue
|
|
if len(tokens) < _OUT_DIR_NAME_MAX_FIELDS:
|
|
tokens.append(token)
|
|
else:
|
|
overflow.append(f"{label}={token}")
|
|
|
|
name = now.strftime("%Y%m%d_%H%M")
|
|
if tokens:
|
|
name += "_" + "_".join(tokens)
|
|
if overflow:
|
|
digest = hashlib.md5("|".join(sorted(overflow)).encode()).hexdigest()[:6]
|
|
name += f"_+{len(overflow)}more-{digest}"
|
|
return name
|
|
|
|
|
|
def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path:
|
|
"""Auto-derived out dir from cfg's hyperparams (see `default_out_dir_name`),
|
|
with a numeric suffix loop so two runs whose name collides (same
|
|
non-default hyperparams, same to-the-minute timestamp) don't clobber each
|
|
other's directory. Shared by `giant train` and `giant new-run`.
|
|
"""
|
|
base_name = default_out_dir_name(cfg)
|
|
out_dir = base / base_name
|
|
suffix = 2
|
|
while out_dir.exists():
|
|
out_dir = base / f"{base_name}_{suffix}"
|
|
suffix += 1
|
|
return out_dir
|
|
|
|
|
|
def seed_everything(seed: int) -> None:
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def _toml_value(v) -> str:
|
|
if isinstance(v, bool):
|
|
return "true" if v else "false"
|
|
if isinstance(v, str):
|
|
return repr(v)
|
|
return str(v)
|
|
|
|
|
|
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)
|
|
|
|
|
|
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:<18} = {_toml_value(v)}")
|
|
|
|
(out_dir / "config.toml").write_text("\n".join(lines))
|
|
|
|
|
|
def build_run_meta(
|
|
data: Path,
|
|
seed: int,
|
|
n_pdg_codes: int,
|
|
n_materials: int,
|
|
n_train_events: int,
|
|
n_val_events: int,
|
|
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"),
|
|
"python_version": sys.version.split()[0],
|
|
"torch_version": torch.__version__,
|
|
"command": " ".join(sys.argv),
|
|
"data_path": str(data),
|
|
"n_pdg_codes": n_pdg_codes,
|
|
"n_materials": n_materials,
|
|
"n_train_events": n_train_events,
|
|
"n_val_events": n_val_events,
|
|
"n_train_steps": n_train_steps,
|
|
}
|