From 72f5a891bf30249c46234c50157efeec522d9543 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 13 Aug 2026 10:21:13 +0200 Subject: [PATCH] Split giant/model/network.py into giant/model/ (issues.md Issue 8) Pure file-move refactor: network.py's 1742 lines held six distinct concerns (layers, condition encoder, routers, trunks, history encoders, stage models, legacy migration, builders) that the v0.3.0 composable-parts refactor already separated at the class level but not the file level. Split along those seams into layers.py/encoders.py/routers.py/trunks.py/ history.py/models.py/_legacy.py/builders.py; network.py is now an 83-line re-export shim so no external import site needed to change. No logic, signature, or behavior changes. --- giant/model/_legacy.py | 114 +++ giant/model/builders.py | 200 +++++ giant/model/encoders.py | 122 +++ giant/model/history.py | 154 ++++ giant/model/layers.py | 76 ++ giant/model/models.py | 574 ++++++++++++ giant/model/network.py | 1819 ++------------------------------------- giant/model/routers.py | 375 ++++++++ giant/model/trunks.py | 155 ++++ issues.md | 25 +- 10 files changed, 1875 insertions(+), 1739 deletions(-) create mode 100644 giant/model/_legacy.py create mode 100644 giant/model/builders.py create mode 100644 giant/model/encoders.py create mode 100644 giant/model/history.py create mode 100644 giant/model/layers.py create mode 100644 giant/model/models.py create mode 100644 giant/model/routers.py create mode 100644 giant/model/trunks.py diff --git a/giant/model/_legacy.py b/giant/model/_legacy.py new file mode 100644 index 0000000..abf3977 --- /dev/null +++ b/giant/model/_legacy.py @@ -0,0 +1,114 @@ +"""v0.2 -> v0.3 checkpoint migration: translates a v0.2 checkpoint's flat +`model_config`/state dicts into the current nested shape (issues.md Issue 8; +see also `giant._migration` and `giant.config.migrate_config`, the sibling +config.toml migration surface — issues.md Issue 6).""" + +from giant._migration import V02_FIXED_FACTS, reject_legacy_router_expert_sizing +from giant.constants import EMB_DIM, K_MAX + + +def _migrate_legacy_model_config(model_config: dict) -> dict: + """Translate a v0.2 checkpoint's flat `model_config` (giant/pipeline.py's + old shape: `hidden_dim`/`n_blocks`/`emb_dim`/`dropout`/`conditioning`/ + `router`/`mode`/... all at one level) into the nested + `{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model", + "stage2_model"}` shape `build_models` expects. + + Sets `stage2_model.n_sec.owner = "stage1"` so the n_sec_head weights a v0.2 + checkpoint carries on its Stage-1 module keep loading there instead of the new + default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's + own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not + just be labeled as such. + + Only the monolithic (non-routed) trunk shape is exercised by the step-2 + migration test; a routed v0.2 checkpoint still builds correctly here + (the router config passes through), but its + state dict isn't covered by `migrate_legacy_state_dict` below. + """ + m = model_config + conditioning_mode = m.get("conditioning", "embedding") + generator = m.get("mode", "flow") + hidden_dim = m.get("hidden_dim", 256) + n_blocks = m.get("n_blocks", 6) + emb_dim = m.get("emb_dim", EMB_DIM) + dropout = m.get("dropout", 0.1) + k_max = m.get("k_max", K_MAX) + noise_dim = m.get("noise_dim", 64) + router_cfg = dict(m.get("router") or {}) + reject_legacy_router_expert_sizing(router_cfg, source="this checkpoint's model_config.router") + router_cfg.setdefault("enabled", False) + + F = V02_FIXED_FACTS + cond_n_layers = F["conditioning.particle.n_layers"] # same fact for both axes + return { + "pdg_vocab": m["pdg_vocab"], + "mat_vocab": m["mat_vocab"], + "conditioning": { + "out_dim": F["conditioning.out_dim"], + "share_stages": False, + "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers}, + "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers}, + }, + "stage1_model": { + "active": F["stage1_model.active"], + "generator": generator, + "hidden_dim": hidden_dim, + "n_res_blocks": n_blocks, + "dropout": dropout, + "flow": {"time_dim": F["stage1_model.flow.time_dim"]}, + "ddpm": {"time_dim": F["stage1_model.ddpm.time_dim"]}, + "wgan": {"noise_dim": noise_dim}, + "router": dict(router_cfg), + }, + "stage2_model": { + "active": F["stage2_model.active"], + "decoder": F["stage2_model.decoder"], + "generator": generator, + "hidden_dim": hidden_dim, + "n_res_blocks": n_blocks, + "dropout": dropout, + "k_max": k_max, + "context_dim": F["stage2_model.context_dim"], + "n_sec": {"mode": "head", "owner": "stage1"}, + "particle_type": {"target": F["stage2_model.particle_type.target"]}, + "flow": {"time_dim": F["stage2_model.flow.time_dim"]}, + "ddpm": {"time_dim": F["stage2_model.ddpm.time_dim"]}, + "wgan": {"noise_dim": noise_dim}, + "router": {**router_cfg, "tie_to_stage1": False}, + }, + } + + +def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple[dict, dict]: + """Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`, + `SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new + `(Stage1Model, Stage2OneShot)` module structure produced by + `build_models(_migrate_legacy_model_config(model_config))`. + + Only the monolithic (non-routed) trunk shape is handled. + """ + + def _trunk_prefix(k: str) -> str: + if k.startswith(("input_proj.", "blocks.", "out_proj.")): + return f"trunk.{k}" + return k + + new_stage1 = {} + for k, v in old_stage1_sd.items(): + if k.startswith("n_sec_head."): + new_stage1[k] = v # stays top-level (n_sec.owner="stage1") + else: + new_stage1[_trunk_prefix(k)] = v + + new_stage2 = {} + for k, v in old_stage2_sd.items(): + if k.startswith("cond_enc.base."): + new_stage2["cond_enc." + k[len("cond_enc.base.") :]] = v + elif k.startswith("cond_enc.stage1_proj."): + new_stage2["context_adapter.proj." + k[len("cond_enc.stage1_proj.") :]] = v + elif k.startswith("cond_enc.fuse."): + new_stage2["fuse." + k[len("cond_enc.fuse.") :]] = v + else: + new_stage2[_trunk_prefix(k)] = v + + return new_stage1, new_stage2 diff --git a/giant/model/builders.py b/giant/model/builders.py new file mode 100644 index 0000000..887a573 --- /dev/null +++ b/giant/model/builders.py @@ -0,0 +1,200 @@ +"""Factories: `build_models`/`build_critics` assemble the top-level stage +models from a config dict (issues.md Issue 8).""" + +import torch.nn as nn + +from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig +from giant.constants import X_DIM +from giant.model._legacy import _migrate_legacy_model_config +from giant.model.encoders import ConditionEncoder +from giant.model.models import CriticModel, Stage1Model, Stage2Autoregressive, Stage2OneShot, stage2_trunk_sec_dim +from giant.model.routers import Router, _build_router_from_cfg + +# --------------------------------------------------------------------------- +# Factories +# --------------------------------------------------------------------------- + + +def build_models(model_config: dict) -> dict[str, nn.Module | None]: + """Construct `{"stage1": ..., "stage2": ...}` from a config dict — either + the new nested shape (has a `"stage1_model"` key, plus `"pdg_vocab"`/ + `"mat_vocab"`/`"conditioning"` at the top level) or a v0.2 checkpoint's + flat `model_config`, auto-migrated via `_migrate_legacy_model_config`. + + A stage is `None` in the result when that stage's `active = False`. + `stage2_model.router.tie_to_stage1` shares stage 1's literal `Router` + instance rather than building a second, independently-parameterized one + (v0.2's actual — probably accidental — behaviour: two routers built from + one config with no semantic relationship between them). + + `conditioning.share_stages = true` builds one `ConditionEncoder` + instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/ + `Stage2Autoregressive`'s `cond_enc` param), instead of each stage + building its own — halving the conditioning parameter count and forcing a + common representation. `false` (default) keeps v0.2 behaviour: + independent instances with identical config but independent weights. + """ + cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config) + pdg_vocab = cfg["pdg_vocab"] + mat_vocab = cfg["mat_vocab"] + conditioning = cfg["conditioning"] + particle_cfg = conditioning["particle"] + material_cfg = conditioning["material"] + particle_conditioning = particle_cfg["type"] + 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_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 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 = 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 + n_sec_owner = s2_spec.n_sec.owner + n_sec_head_k_max = s2_spec.k_max if n_sec_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=s1_spec.hidden_dim, + n_res_blocks=s1_spec.n_res_blocks, + cond_out_dim=cond_out_dim, + dropout=s1_spec.dropout, + generator=generator, + 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 s2_spec.active: + decoder = s2_spec.decoder + router_cfg = cfg["stage2_model"].get("router") or {} + stage2_router: Router | None = 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 = 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 + n_sec_owner = s2_spec.n_sec.owner + k_max = s2_spec.k_max + particle_type_cfg = s2_spec.particle_type.to_dict() + + if decoder == "autoregressive": + 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=s2_spec.hidden_dim, + n_res_blocks=s2_spec.n_res_blocks, + cond_out_dim=cond_out_dim, + context_dim=s2_spec.context_dim, + dropout=s2_spec.dropout, + generator=generator, + time_dim=time_dim, + noise_dim=s2_spec.wgan.noise_dim, + k_max=k_max, + router=stage2_router, + build_n_sec_head=n_sec_owner != "stage1", + particle_type_cfg=particle_type_cfg, + 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: + sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, k_max, particle_cfg["emb_dim"]) + result["stage2"] = Stage2OneShot( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=s2_spec.hidden_dim, + n_res_blocks=s2_spec.n_res_blocks, + cond_out_dim=cond_out_dim, + context_dim=s2_spec.context_dim, + sec_dim=sec_dim, + dropout=s2_spec.dropout, + generator=generator, + time_dim=time_dim, + noise_dim=s2_spec.wgan.noise_dim, + k_max=k_max, + router=stage2_router, + build_n_sec_head=n_sec_owner != "stage1", + particle_type_cfg=particle_type_cfg, + cond_enc=shared_cond_enc, + ) + + return result + + +def build_critics(model_config: dict) -> dict[str, nn.Module | None]: + """Construct `{"stage1": ..., "stage2": ...}` critics for `generator = + "wgan"` training. Training-only — never persisted for inference the way + `build_models`'s pair is. `None` for a stage that's inactive or not + WGAN.""" + cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config) + pdg_vocab = cfg["pdg_vocab"] + mat_vocab = cfg["mat_vocab"] + conditioning = cfg["conditioning"] + particle_cfg = conditioning["particle"] + material_cfg = conditioning["material"] + 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 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=s1_spec.hidden_dim, + n_res_blocks=s1_spec.n_res_blocks, + cond_out_dim=cond_out_dim, + dropout=s1_spec.dropout, + stage="stage1", + ) + + 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, + mat_vocab=mat_vocab, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + in_dim=in_dim, + hidden_dim=s2_spec.hidden_dim, + n_res_blocks=s2_spec.n_res_blocks, + cond_out_dim=cond_out_dim, + dropout=s2_spec.dropout, + stage="stage2", + context_dim=s2_spec.context_dim, + ) + + return result diff --git a/giant/model/encoders.py b/giant/model/encoders.py new file mode 100644 index 0000000..384903b --- /dev/null +++ b/giant/model/encoders.py @@ -0,0 +1,122 @@ +"""Conditioning encoder — fuses continuous conditioning with particle/material +identity (issues.md Issue 8).""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM +from giant.model.layers import _make_axis_mlp + + +def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]: + """`cond_cat` column indices for each axis's top-N-onehot index, or + `None` if that axis isn't `"onehot"`. + + Columns 0/1 are always the dense pdg/material vocab index. The particle + top-N column (if any) comes next, then the material top-N column (if + any) — `giant.data.transforms.build_cond_features`/`build_features` + append columns in this same order, so the two sides must never drift + apart. + """ + col = 2 + particle_col = None + if particle_type == "onehot": + particle_col = col + col += 1 + material_col = None + if material_type == "onehot": + material_col = col + col += 1 + return particle_col, material_col + + +class ConditionEncoder(nn.Module): + """Fuses continuous conditioning with particle/material identity. + + The particle and material axes are configured independently + (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`) + and may mix freely, e.g. material "physical" with particle "embedding". + Three modes per axis: + - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s + dense training-vocab index. Memorizes the training menu. + - "physical": an `n_layers`-deep MLP over the axis's raw physical + properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see + giant.data.transforms.build_features), computable for any PDG code / + material name rather than only ones seen in training. + - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other + class map (`giant.data.loader.build_topn_map_from_files`/ + `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra + top-N-index column(s) — see `_cat_col_layout`. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + cont_dim: int = COND_DIM, + out_dim: int = 128, + ) -> None: + super().__init__() + self.particle_cfg = dict(particle_cfg) + self.material_cfg = dict(material_cfg) + self._particle_topn_col, self._material_topn_col = cat_col_layout(particle_cfg["type"], material_cfg["type"]) + + p_type = particle_cfg["type"] + p_emb_dim = particle_cfg["emb_dim"] + if p_type == "embedding": + self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim) + elif p_type == "physical": + self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1)) + elif p_type != "onehot": + raise ValueError(f"unknown conditioning.particle.type {p_type!r}") + + m_type = material_cfg["type"] + m_emb_dim = material_cfg["emb_dim"] + if m_type == "embedding": + self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim) + elif m_type == "physical": + self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1)) + elif m_type != "onehot": + raise ValueError(f"unknown conditioning.material.type {m_type!r}") + + in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim + self.mlp = nn.Sequential( + nn.Linear(in_dim, out_dim), + nn.SiLU(), + nn.Linear(out_dim, out_dim), + ) + + def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): + p_type = self.particle_cfg["type"] + if p_type == "embedding": + return self.pdg_emb(cond_cat[:, 0]) + if p_type == "physical": + particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM] + return self.particle_mlp(particle_phys) + assert self._particle_topn_col is not None + return F.one_hot( + cond_cat[:, self._particle_topn_col], + num_classes=self.particle_cfg["emb_dim"], + ).float() + + def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): + m_type = self.material_cfg["type"] + if m_type == "embedding": + return self.mat_emb(cond_cat[:, 1]) + if m_type == "physical": + material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :] + return self.material_mlp(material_phys) + assert self._material_topn_col is not None + return F.one_hot( + cond_cat[:, self._material_topn_col], + num_classes=self.material_cfg["emb_dim"], + ).float() + + def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + pdg_e = self._particle_embed(cond_cont, cond_cat) + mat_e = self._material_embed(cond_cont, cond_cat) + x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1) + return self.mlp(x) diff --git a/giant/model/history.py b/giant/model/history.py new file mode 100644 index 0000000..80fd22b --- /dev/null +++ b/giant/model/history.py @@ -0,0 +1,154 @@ +"""History encoders — stage-2 autoregressive only. Self-contained, no +dependency on any other `giant.model` submodule (issues.md Issue 8).""" + +import torch +import torch.nn as nn + + +class HistoryEncoder(nn.Module): + """Interface for stage-2 autoregressive per-token history summaries: + `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over + a full (teacher-forced) token sequence — used by training. `MarkovHistory` + and `AttentionHistory` are the two implementations. Inference + (`giant/sample.py`) generates one token at a + time and cannot afford `forward`'s per-step cost to be O(K) (attention + would then be O(K^2) over a rollout's k_max loop); encoders that need + incremental state for that path additionally implement `init_cache`/ + `step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since + its per-step cost is already O(1) (it only ever looks at the previous + token, not the full prefix).""" + + def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + +class MarkovHistory(HistoryEncoder): + """Summarizes the previous secondary's own `(energy_fraction, direction, + type_representation)` through one small MLP — the "markov" history: + token i+1 only ever sees token i plus the running scalars + (`remaining_frac`/`slot_idx`, fused in separately by + `Stage2Autoregressive._token_cond`), not the full prefix. + + At slot 0 (`has_prev` False) substitutes a learned start vector rather + than zeros — a reasonable default. + """ + + def __init__(self, in_dim: int, out_dim: int) -> None: + super().__init__() + self.start = nn.Parameter(torch.zeros(in_dim)) + self.mlp = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU()) + + def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: + start = self.start.view(1, 1, -1).expand_as(feat) + x = torch.where(has_prev.unsqueeze(-1), feat, start) + return self.mlp(x) + + +class _CausalAttnBlock(nn.Module): + """One pre-norm causal self-attention block for `AttentionHistory`. + + Exposes two forward paths that must agree (see + `test_attention_history_step_matches_forward` in `tests/test_network.py`): + `forward` — the full-sequence, causally-masked pass used for training; + `step` — an incremental pass for inference, given the *pre-attention* + normalized hidden states of every earlier position (`kv_cache`, i.e. + `norm1(x)` for positions `< t`, not `x` itself). Caching `norm1(x)` rather + than raw `x` is what makes `step` correct: this block's attention needs + exactly that quantity as keys/values, and `LayerNorm` has no cross-position + interaction, so recomputing it per position instead of caching it would + still be correct but pointlessly repeat work. The *next* block's cache is + built from a different sequence (this block's output), so each block owns + an independent cache entry. + """ + + def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True) + self.norm2 = nn.LayerNorm(dim) + self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)) + + def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor: + h = self.norm1(x) + attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False) + x = x + attn_out + x = x + self.mlp(self.norm2(x)) + return x + + def step(self, x_new: torch.Tensor, kv_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]: + """`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None` + (first position) or `(B, T, dim)` — `norm1(x)` of every earlier + position at this same block. Returns `(out, new_kv_cache)`, `out` + being this position's block output (`(B, 1, dim)`, to feed the next + block's `step`), `new_kv_cache` the same cache extended by this + position (to reuse at this block's *next* `step` call).""" + h_new = self.norm1(x_new) + kv = h_new if kv_cache is None else torch.cat([kv_cache, h_new], dim=1) + attn_out, _ = self.attn(h_new, kv, kv, need_weights=False) + x = x_new + attn_out + x = x + self.mlp(self.norm2(x)) + return x, kv + + +class AttentionHistory(HistoryEncoder): + """Causal self-attention over the emitted-token prefix — the more + expressive alternative to `MarkovHistory`'s fixed previous-token-only + summary. `feat`/`has_prev` + follow the same shifted-by-one convention `MarkovHistory` and + `Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s + own `(energy_fraction, direction, type_representation)`, with a learned + start vector substituted at `has_prev == False` positions (only slot 0 in + practice — see `giant.training.stage2_inputs._ar_has_prev`). Causal masking then makes + position `i`'s output a function of `feat[:, 1:i+1]` — i.e. tokens + `0..i-1` — exactly the prefix available when predicting token `i`. + + `forward` is the parallel training path (one pass over the whole + teacher-forced sequence); `init_cache`/`step` are the incremental + inference path `giant/sample.py` uses, one new token per call, to avoid + re-encoding the whole prefix from scratch every slot — `step` must be + called exactly once per slot (its cache-extension is not idempotent), + so a slot's output must be reused for + every model call within that slot (`forward`'s ODE substeps, or a separate + `predict_type` call) rather than re-derived — see + `Stage2Autoregressive.history_step`. + """ + + def __init__(self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2) -> None: + super().__init__() + self.start = nn.Parameter(torch.zeros(in_dim)) + self.in_proj = nn.Linear(in_dim, out_dim) + self.blocks = nn.ModuleList([_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)]) + + def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: + start = self.start.view(1, 1, -1).expand_as(feat) + x = torch.where(has_prev.unsqueeze(-1), feat, start) + return self.in_proj(x) + + def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: + B, K, _ = feat.shape + x = self._embed(feat, has_prev) + mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device) + for block in self.blocks: + x = block(x, mask) + return x + + def init_cache(self) -> list[torch.Tensor | None]: + return [None for _ in self.blocks] + + def step( + self, + token_feat: torch.Tensor, + has_prev: torch.Tensor, + cache: list[torch.Tensor | None], + ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: + """`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest + token's own features (what would be `feat[:, k]` in `forward`). + Advances every block's cache by this position and returns this + position's output (`(B, 1, out_dim)`, the correct history summary for + the NEXT slot) plus the updated cache.""" + x = self._embed(token_feat, has_prev) + new_cache: list[torch.Tensor | None] = [] + for block, kv in zip(self.blocks, cache): + x, kv_new = block.step(x, kv) + new_cache.append(kv_new) + return x, new_cache diff --git a/giant/model/layers.py b/giant/model/layers.py new file mode 100644 index 0000000..84834b1 --- /dev/null +++ b/giant/model/layers.py @@ -0,0 +1,76 @@ +"""Small stateless-ish building blocks shared across encoders/trunks/models — +no dependency on any other `giant.model` submodule (issues.md Issue 8).""" + +import math + +import torch +import torch.nn as nn + + +class SinusoidalEmbedding(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + assert dim % 2 == 0, "dim must be even" + half = dim // 2 + freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1)) + self.register_buffer("freqs", freqs) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + t = t.reshape(-1, 1).float() + args = t * self.freqs.unsqueeze(0) # (B, half) + return torch.cat([args.sin(), args.cos()], dim=-1) # (B, dim) + + +def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential: + """`n_layers`-deep MLP producing an `emb_dim`-wide vector from `in_dim` + physical properties (`conditioning.{particle,material}.n_layers`). + + `n_layers=1` (the v0.3.0 default): a single `Linear`, no hidden + activation. `n_layers=2` reproduces v0.2's hardcoded depth exactly — + `Linear -> SiLU -> Linear` — which is why `migrate_config` back-fills + `n_layers=2` for migrated configs rather than the v0.3 default of 1 (see + its docstring). + """ + if n_layers < 1: + raise ValueError(f"n_layers must be >= 1, got {n_layers}") + if n_layers == 1: + return nn.Sequential(nn.Linear(in_dim, emb_dim)) + layers: list[nn.Module] = [nn.Linear(in_dim, emb_dim), nn.SiLU()] + for _ in range(n_layers - 2): + layers += [nn.Linear(emb_dim, emb_dim), nn.SiLU()] + layers.append(nn.Linear(emb_dim, emb_dim)) + return nn.Sequential(*layers) + + +class ContextAdapter(nn.Module): + """Projects a stage's outcome (e.g. Stage 1's 9D target) down to a + fixed-width context vector for a downstream stage's conditioning — + `stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj` + (+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since + `SecondaryConditionEncoder` as a wrapper class disappears.""" + + def __init__(self, in_dim: int, context_dim: int) -> None: + super().__init__() + self.proj = nn.Linear(in_dim, context_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.tanh(self.proj(x)) + + +class ResBlock(nn.Module): + def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None: + super().__init__() + self.norm = nn.LayerNorm(dim) + self.linear1 = nn.Linear(dim, dim) + self.cond_proj = nn.Linear(cond_dim, dim, bias=False) + self.act = nn.SiLU() + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim, dim) + + def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + h = self.norm(x) + h = self.linear1(h) + self.cond_proj(cond) + h = self.act(h) + h = self.dropout(h) + h = self.linear2(h) + return x + h diff --git a/giant/model/models.py b/giant/model/models.py new file mode 100644 index 0000000..a57d976 --- /dev/null +++ b/giant/model/models.py @@ -0,0 +1,574 @@ +"""Top-level stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`, +`CriticModel` — composed from encoders/trunks/history (issues.md Issue 8).""" + +import torch +import torch.nn as nn + +from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM +from giant.model.encoders import ConditionEncoder +from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory +from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding +from giant.model.routers import Router +from giant.model.trunks import build_trunk + +# --------------------------------------------------------------------------- +# Stage models +# --------------------------------------------------------------------------- + + +def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int: + """Width of a single secondary slot's type slice — + `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else + `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are + `conditioning.particle.emb_dim` wide).""" + target = particle_type_cfg.get("target", "physical") + return PARTICLE_PHYS_DIM if target == "physical" else emb_dim + + +def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int: + """`Stage2OneShot`'s trunk output width. + + `target = "physical"` is untouched from v0.2/today: + `k_max * SEC_SLOT_DIM`, the type slice folded into the same + flow-matched/WGAN vector as the continuous stick/dir slots. + + `target` in `("onehot", "embedding")`: under `generator == "wgan"` the + type slice is still folded in (adversarial for onehot via ST-Gumbel, + already-continuous for embedding), just `emb_dim` wide instead of + `PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under + `generator in ("flow", "ddpm")` the type slice isn't part of this vector + at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk + only covers `k_max * CONT_SLOT_DIM`. + """ + target = particle_type_cfg.get("target", "physical") + if target == "physical": + return k_max * SEC_SLOT_DIM + if generator == "wgan": + return k_max * (CONT_SLOT_DIM + emb_dim) + return k_max * CONT_SLOT_DIM + + +class Stage1Model(nn.Module): + """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs + move it to stage 2, except for a migrated v0.2 checkpoint + (`n_sec_head_k_max` given), where it stays attached here + since that's where its weights live and what conditioning it was trained + against (see `_migrate_legacy_model_config`). + + `cond_enc`, if given, is used in place of building a fresh + `ConditionEncoder` — `conditioning.share_stages = true`: `build_models` + constructs one shared instance and passes it to both stages, halving the + conditioning parameter count and forcing a common representation.""" + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + hidden_dim: int = 256, + n_res_blocks: int = 6, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + dropout: float = 0.0, + generator: str = "flow", + time_dim: int = 64, + noise_dim: int = 64, + router: Router | None = None, + n_sec_head_k_max: int | None = None, + cond_enc: ConditionEncoder | None = None, + ) -> None: + super().__init__() + self.generator_kind = generator + self.noise_dim = noise_dim + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + ) + has_time = generator in ("flow", "ddpm") + self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None + merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim + in_dim = noise_dim if generator == "wgan" else x_dim + self.trunk = build_trunk(router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout) + self.n_sec_head = None + if n_sec_head_k_max is not None: + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, n_sec_head_k_max + 1), + ) + + def forward( + self, + x_t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + t: torch.Tensor | None = None, + ) -> torch.Tensor: + c_emb = self.cond_enc(cond_cont, cond_cat) + cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb + return self.trunk(x_t, cond, cond_cont, cond_cat) + + def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only + valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0 + configs predict n_sec from Stage2OneShot instead.""" + if self.n_sec_head is None: + raise RuntimeError( + "this Stage1Model has no n_sec_head — n_sec now lives on " + "stage 2 by default; this method only exists " + "for a migrated v0.2 checkpoint (n_sec.owner='stage1')" + ) + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + +class Stage2OneShot(nn.Module): + """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour, + reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`, + step 4/5, not implemented yet). + + Owns `n_sec_head` by default unless `build_n_sec_head=False` + (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to + Stage1Model — see `_migrate_legacy_model_config`). + + `particle_type_cfg["target"]` (default `"physical"`) selects the + secondary-type mechanism: `"physical"` keeps the type slice folded into + the trunk's own + flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by + the caller via `stage2_trunk_sec_dim` — already reflects this). Under + `"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type + slice is predicted by a separate `type_head` instead (same shape pattern + as `n_sec_head`) — `sec_dim` then covers only the continuous + stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors. + Under `generator == "wgan"` the type slice stays folded into `sec_dim` + (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is + unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation. + + `cond_enc`, if given, is used in place of building a fresh + `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + hidden_dim: int = 256, + n_res_blocks: int = 6, + cond_out_dim: int = 128, + context_dim: int = 64, + sec_dim: int = SEC_DIM, + x_dim: int = X_DIM, + dropout: float = 0.0, + generator: str = "wgan", + time_dim: int = 64, + noise_dim: int = 64, + k_max: int = K_MAX, + router: Router | None = None, + build_n_sec_head: bool = True, + particle_type_cfg: dict | None = None, + cond_enc: ConditionEncoder | None = None, + ) -> None: + super().__init__() + self.generator_kind = generator + self.noise_dim = noise_dim + self.k_max = k_max + self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) + self.type_dim = stage2_type_dim(self.particle_type_cfg, particle_cfg["emb_dim"]) + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + ) + self.context_adapter = ContextAdapter(x_dim, context_dim) + self.fuse = nn.Sequential( + nn.Linear(cond_out_dim + context_dim, cond_out_dim), + nn.SiLU(), + ) + has_time = generator in ("flow", "ddpm") + self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None + merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim + in_dim = noise_dim if generator == "wgan" else sec_dim + self.trunk = build_trunk(router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout) + self.n_sec_head = None + if build_n_sec_head: + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, k_max + 1), + ) + self.type_head = None + target = self.particle_type_cfg.get("target", "physical") + if target != "physical" and generator in ("flow", "ddpm"): + emb_dim = particle_cfg["emb_dim"] + self.type_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, k_max * emb_dim), + ) + self._type_k_max = k_max + self._type_emb_dim = emb_dim + + def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: + base = self.cond_enc(cond_cont, cond_cat) + ctx = self.context_adapter(stage1_out) + return self.fuse(torch.cat([base, ctx], dim=-1)) + + def forward( + self, + x_t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + t: torch.Tensor | None = None, + ) -> torch.Tensor: + c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) + cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb + return self.trunk(x_t, cond, cond_cont, cond_cat) + + def predict_n_sec( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + if self.n_sec_head is None: + raise RuntimeError( + "this Stage2OneShot has no n_sec_head — it belongs to a " + "migrated v0.2 checkpoint (n_sec.owner='stage1'); call " + "stage1.predict_n_sec(cond_cont, cond_cat) instead" + ) + c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) + return self.n_sec_head(c_emb) + + def predict_type( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + """`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or + vectors (`target="embedding"`) — only under `generator in ("flow", + "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s + own output instead (see class docstring).""" + if self.type_head is None: + raise RuntimeError( + "this Stage2OneShot has no type_head — either " + "particle_type.target='physical' (the type slice is part of " + "forward()'s own output) or generator='wgan' (the WGAN " + "trainer reads the type slice out of forward()'s output " + "directly instead)" + ) + c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) + return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim) + + +class Stage2Autoregressive(nn.Module): + """Emits secondaries one at a time in descending-energy order, instead + of `Stage2OneShot`'s simultaneous + k_max-slot prediction. `history` selects `MarkovHistory` or + `AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only). + `teacher_forcing` handling lives entirely in the trainer + (`giant/train.py`), since it only affects how training inputs are + assembled, not this module's architecture. + + Under teacher forcing every token's conditioning is built from ground + truth, so a whole K-token sequence trains in one parallel batched pass: + `forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded + to `k_max`) — this also means a future one-token-at-a-time inference loop + (`K=1` per call, step 6) needs no interface change here. + + Two independent conditioning paths, mirroring `Stage2OneShot`'s + `_cond_embed` but split in two: `_base_cond` (`cond_enc` + + `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend + on token position; `_token_cond` additionally fuses in the history + encoding and two running scalars (remaining energy-budget fraction, + normalized slot index), and feeds `forward`/`predict_type`/the trunk. + + `cond_enc`, if given, is used in place of building a fresh + `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + hidden_dim: int = 256, + n_res_blocks: int = 6, + cond_out_dim: int = 128, + context_dim: int = 64, + x_dim: int = X_DIM, + dropout: float = 0.0, + generator: str = "wgan", + time_dim: int = 64, + noise_dim: int = 64, + k_max: int = K_MAX, + router: Router | None = None, + build_n_sec_head: bool = True, + particle_type_cfg: dict | None = None, + history: str = "markov", + attn_n_heads: int = 4, + attn_n_layers: int = 2, + cond_enc: ConditionEncoder | None = None, + ) -> None: + super().__init__() + if history not in ("markov", "attention"): + raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'") + self.history_kind = history + self.generator_kind = generator + self.noise_dim = noise_dim + self.k_max = k_max + self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) + emb_dim = particle_cfg["emb_dim"] + self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim) + + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + ) + self.context_adapter = ContextAdapter(x_dim, context_dim) + self.base_fuse = nn.Sequential( + nn.Linear(cond_out_dim + context_dim, cond_out_dim), + nn.SiLU(), + ) + + # Reuses conditioning.out_dim for the history encoder's own output + # width — there's no dedicated stage2_model.autoregressive key for + # this, a reasonable default rather than a design-doc-specified value. + history_dim = cond_out_dim + hist_in_dim = CONT_SLOT_DIM + self.type_dim + self.history_encoder: HistoryEncoder = ( + AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers) + if history == "attention" + else MarkovHistory(hist_in_dim, history_dim) + ) + token_fuse_in = cond_out_dim + context_dim + history_dim + 2 # +2: remaining_frac, slot_idx + self.token_fuse = nn.Sequential( + nn.Linear(token_fuse_in, cond_out_dim), + nn.SiLU(), + ) + + has_time = generator in ("flow", "ddpm") + self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None + merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim + token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim) + in_dim = noise_dim if generator == "wgan" else token_dim + self.trunk = build_trunk( + router, + in_dim, + token_dim, + hidden_dim, + n_res_blocks, + merged_cond_dim, + dropout, + ) + + self.n_sec_head = None + if build_n_sec_head: + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, k_max + 1), + ) + self.type_head = None + target = self.particle_type_cfg.get("target", "physical") + if target != "physical" and generator in ("flow", "ddpm"): + self.type_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, self.type_dim), + ) + + def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: + base = self.cond_enc(cond_cont, cond_cat) + ctx = self.context_adapter(stage1_out) + return self.base_fuse(torch.cat([base, ctx], dim=-1)) + + def _token_cond( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + hist: torch.Tensor | None = None, + ) -> torch.Tensor: + """`hist`, if given, overrides recomputing `self.history_encoder` + from `history_feat`/`has_prev` — the inference-time KV-cache path + (`Stage2Autoregressive.history_step`) precomputes it once per slot and + passes it in here so a slot's (possibly several) model calls — an ODE + loop's substeps, or a separate `predict_type` call — read the same + cached history instead of each re-deriving (and, under attention, + re-appending to the cache — see `AttentionHistory.step`'s docstring).""" + K = history_feat.size(1) + base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1) + ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1) + if hist is None: + hist = self.history_encoder(history_feat, has_prev) + scalars = torch.stack([remaining_frac, slot_idx], dim=-1) + return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1)) + + def init_history_cache(self): + """Inference-only incremental-decoding state for `self.history_encoder` + (`giant/sample.py`'s AR loop): `None` under `history="markov"` (its + per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or + `AttentionHistory.init_cache()` under `history="attention"`.""" + if isinstance(self.history_encoder, AttentionHistory): + return self.history_encoder.init_cache() + return None + + def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]: + """One inference slot's worth of history encoding: advances `cache` + (from `init_history_cache`, or a previous `history_step` call) by + `token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous + token, same convention `giant.sample.sample_secondaries_ar` already + threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is + this slot's history summary (pass it as `_token_cond`'s `hist=` to + every model call made for this slot), `new_cache` is what to pass into + the *next* slot's `history_step`. Must be called exactly once per + slot — see `AttentionHistory.step`'s docstring.""" + if isinstance(self.history_encoder, AttentionHistory): + return self.history_encoder.step(token_feat, has_prev, cache) + return self.history_encoder(token_feat, has_prev), cache + + def forward( + self, + x_t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + t: torch.Tensor | None = None, + hist: torch.Tensor | None = None, + ) -> torch.Tensor: + B, K = x_t.shape[0], x_t.shape[1] + c_emb = self._token_cond( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + hist=hist, + ) + if self.time_emb is not None: + assert t is not None + t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1) + cond = torch.cat([t_emb, c_emb], dim=-1) + else: + cond = c_emb + x_flat = x_t.reshape(B * K, -1) + cond_flat = cond.reshape(B * K, -1) + cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1) + cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1) + out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat) + return out.view(B, K, -1) + + def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: + if self.n_sec_head is None: + raise RuntimeError( + "this Stage2Autoregressive has no n_sec_head — it belongs to " + "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call " + "stage1.predict_n_sec(cond_cont, cond_cat) instead" + ) + return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out)) + + def predict_type( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + history_feat: torch.Tensor, + has_prev: torch.Tensor, + remaining_frac: torch.Tensor, + slot_idx: torch.Tensor, + hist: torch.Tensor | None = None, + ) -> torch.Tensor: + if self.type_head is None: + raise RuntimeError( + "this Stage2Autoregressive has no type_head — either " + "particle_type.target='physical' (the type slice is part of " + "forward()'s own output) or generator='wgan' (the WGAN " + "trainer reads the type slice out of forward()'s output " + "directly instead)" + ) + c_emb = self._token_cond( + cond_cont, + cond_cat, + stage1_out, + history_feat, + has_prev, + remaining_frac, + slot_idx, + hist=hist, + ) + B, K, _ = c_emb.shape + return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim) + + +class CriticModel(nn.Module): + """Generator-agnostic WGAN-GP critic body: a scalar realism score, for + either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"` + mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as + `Stage2OneShot`). Used only when that stage's `generator == "wgan"`.""" + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + particle_cfg: dict, + material_cfg: dict, + in_dim: int, + hidden_dim: int = 256, + n_res_blocks: int = 6, + cond_out_dim: int = 128, + dropout: float = 0.0, + stage: str = "stage1", + context_dim: int = 64, + context_in_dim: int = X_DIM, + ) -> None: + super().__init__() + if stage not in ("stage1", "stage2"): + raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}") + self.stage = stage + self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) + if stage == "stage2": + self.context_adapter = ContextAdapter(context_in_dim, context_dim) + self.fuse = nn.Sequential( + nn.Linear(cond_out_dim + context_dim, cond_out_dim), + nn.SiLU(), + ) + self.input_proj = nn.Linear(in_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)]) + self.out_norm = nn.LayerNorm(hidden_dim) + self.out_proj = nn.Linear(hidden_dim, 1) + + def forward( + self, + x: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor | None = None, + ) -> torch.Tensor: + base = self.cond_enc(cond_cont, cond_cat) + if self.stage == "stage2": + ctx = self.context_adapter(stage1_out) + cond = self.fuse(torch.cat([base, ctx], dim=-1)) + else: + cond = base + h = self.input_proj(x) + for block in self.blocks: + h = block(h, cond) + return self.out_proj(self.out_norm(h)).squeeze(-1) diff --git a/giant/model/network.py b/giant/model/network.py index 4289529..c902c58 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -1,1742 +1,85 @@ -import inspect -import math -import re -from collections.abc import Sequence +"""Re-export shim. -import torch -import torch.nn as nn -import torch.nn.functional as F +`network.py` used to hold every network-related class in one 1742-line file. +It's now split by concern into `giant/model/{layers,encoders,routers,trunks, +history,models,_legacy,builders}.py` (issues.md Issue 8); this module just +re-exports the public surface so every existing `from giant.model.network +import X` call site keeps working unchanged. +""" -from giant._migration import V02_FIXED_FACTS, reject_legacy_router_expert_sizing -from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig -from giant.constants import ( - COND_DIM, - COND_DIM_BASE, - CONT_SLOT_DIM, - EMB_DIM, - K_MAX, - MATERIAL_PHYS_DIM, - PARTICLE_PHYS_DIM, - SEC_DIM, - SEC_SLOT_DIM, - X_DIM, +from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_state_dict +from giant.model.builders import build_critics, build_models +from giant.model.encoders import ConditionEncoder, cat_col_layout +from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory, _CausalAttnBlock +from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, _make_axis_mlp +from giant.model.models import ( + CriticModel, + Stage1Model, + Stage2Autoregressive, + Stage2OneShot, + stage2_trunk_sec_dim, + stage2_type_dim, +) +from giant.model.routers import ( + ROUTER_REGISTRY, + ComposedRouter, + EnergyRouter, + PdgRouter, + ProcessRouter, + Router, + _build_router_from_cfg, + _check_router_conditioning_compat, + _parse_composed_axes, + build_composed_router, + build_router, + register_router, +) +from giant.model.trunks import ( + ExpertTrunk, + MonolithicTrunk, + RoutedTrunk, + Trunk, + _route_forward, + build_trunk, ) -# --------------------------------------------------------------------------- -# Building blocks -# --------------------------------------------------------------------------- - - -class SinusoidalEmbedding(nn.Module): - def __init__(self, dim: int) -> None: - super().__init__() - assert dim % 2 == 0, "dim must be even" - half = dim // 2 - freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1)) - self.register_buffer("freqs", freqs) - - def forward(self, t: torch.Tensor) -> torch.Tensor: - t = t.reshape(-1, 1).float() - args = t * self.freqs.unsqueeze(0) # (B, half) - return torch.cat([args.sin(), args.cos()], dim=-1) # (B, dim) - - -def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]: - """`cond_cat` column indices for each axis's top-N-onehot index, or - `None` if that axis isn't `"onehot"`. - - Columns 0/1 are always the dense pdg/material vocab index. The particle - top-N column (if any) comes next, then the material top-N column (if - any) — `giant.data.transforms.build_cond_features`/`build_features` - append columns in this same order, so the two sides must never drift - apart. - """ - col = 2 - particle_col = None - if particle_type == "onehot": - particle_col = col - col += 1 - material_col = None - if material_type == "onehot": - material_col = col - col += 1 - return particle_col, material_col - - -def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential: - """`n_layers`-deep MLP producing an `emb_dim`-wide vector from `in_dim` - physical properties (`conditioning.{particle,material}.n_layers`). - - `n_layers=1` (the v0.3.0 default): a single `Linear`, no hidden - activation. `n_layers=2` reproduces v0.2's hardcoded depth exactly — - `Linear -> SiLU -> Linear` — which is why `migrate_config` back-fills - `n_layers=2` for migrated configs rather than the v0.3 default of 1 (see - its docstring). - """ - if n_layers < 1: - raise ValueError(f"n_layers must be >= 1, got {n_layers}") - if n_layers == 1: - return nn.Sequential(nn.Linear(in_dim, emb_dim)) - layers: list[nn.Module] = [nn.Linear(in_dim, emb_dim), nn.SiLU()] - for _ in range(n_layers - 2): - layers += [nn.Linear(emb_dim, emb_dim), nn.SiLU()] - layers.append(nn.Linear(emb_dim, emb_dim)) - return nn.Sequential(*layers) - - -class ConditionEncoder(nn.Module): - """Fuses continuous conditioning with particle/material identity. - - The particle and material axes are configured independently - (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`) - and may mix freely, e.g. material "physical" with particle "embedding". - Three modes per axis: - - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s - dense training-vocab index. Memorizes the training menu. - - "physical": an `n_layers`-deep MLP over the axis's raw physical - properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see - giant.data.transforms.build_features), computable for any PDG code / - material name rather than only ones seen in training. - - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other - class map (`giant.data.loader.build_topn_map_from_files`/ - `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra - top-N-index column(s) — see `_cat_col_layout`. - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, - cont_dim: int = COND_DIM, - out_dim: int = 128, - ) -> None: - super().__init__() - self.particle_cfg = dict(particle_cfg) - self.material_cfg = dict(material_cfg) - self._particle_topn_col, self._material_topn_col = cat_col_layout(particle_cfg["type"], material_cfg["type"]) - - p_type = particle_cfg["type"] - p_emb_dim = particle_cfg["emb_dim"] - if p_type == "embedding": - self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim) - elif p_type == "physical": - self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1)) - elif p_type != "onehot": - raise ValueError(f"unknown conditioning.particle.type {p_type!r}") - - m_type = material_cfg["type"] - m_emb_dim = material_cfg["emb_dim"] - if m_type == "embedding": - self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim) - elif m_type == "physical": - self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1)) - elif m_type != "onehot": - raise ValueError(f"unknown conditioning.material.type {m_type!r}") - - in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim - self.mlp = nn.Sequential( - nn.Linear(in_dim, out_dim), - nn.SiLU(), - nn.Linear(out_dim, out_dim), - ) - - def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): - p_type = self.particle_cfg["type"] - if p_type == "embedding": - return self.pdg_emb(cond_cat[:, 0]) - if p_type == "physical": - particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM] - return self.particle_mlp(particle_phys) - assert self._particle_topn_col is not None - return F.one_hot( - cond_cat[:, self._particle_topn_col], - num_classes=self.particle_cfg["emb_dim"], - ).float() - - def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): - m_type = self.material_cfg["type"] - if m_type == "embedding": - return self.mat_emb(cond_cat[:, 1]) - if m_type == "physical": - material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :] - return self.material_mlp(material_phys) - assert self._material_topn_col is not None - return F.one_hot( - cond_cat[:, self._material_topn_col], - num_classes=self.material_cfg["emb_dim"], - ).float() - - def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - pdg_e = self._particle_embed(cond_cont, cond_cat) - mat_e = self._material_embed(cond_cont, cond_cat) - x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1) - return self.mlp(x) - - -class ContextAdapter(nn.Module): - """Projects a stage's outcome (e.g. Stage 1's 9D target) down to a - fixed-width context vector for a downstream stage's conditioning — - `stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj` - (+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since - `SecondaryConditionEncoder` as a wrapper class disappears.""" - - def __init__(self, in_dim: int, context_dim: int) -> None: - super().__init__() - self.proj = nn.Linear(in_dim, context_dim) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return torch.tanh(self.proj(x)) - - -class ResBlock(nn.Module): - def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None: - super().__init__() - self.norm = nn.LayerNorm(dim) - self.linear1 = nn.Linear(dim, dim) - self.cond_proj = nn.Linear(cond_dim, dim, bias=False) - self.act = nn.SiLU() - self.dropout = nn.Dropout(dropout) - self.linear2 = nn.Linear(dim, dim) - - def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: - h = self.norm(x) - h = self.linear1(h) + self.cond_proj(cond) - h = self.act(h) - h = self.dropout(h) - h = self.linear2(h) - return x + h - - -# --------------------------------------------------------------------------- -# Routers — carried over unchanged from v0.2 -# --------------------------------------------------------------------------- - - -class Router(nn.Module): - """Contract for a pluggable mixture-of-experts routing axis. - - Subclasses implement `gate` (soft partition-of-unity weights over - experts, used in train mode for a fully differentiable mixture); - `top1` and `balance_loss` have working defaults so a new routing axis - is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`. - """ - - def __init__(self, n_experts: int) -> None: - super().__init__() - self.n_experts = n_experts - self.gumbel = False - self.gumbel_tau = 1.0 - - def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """(B, n_experts) soft weights, rows summing to 1.""" - raise NotImplementedError - - def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """(B, n_experts) train-time expert-combination weights. - - Default (`gumbel=False`): identical to `gate()`. Opt-in - straight-through Gumbel-softmax (`gumbel=True`, train mode only): - hardens the forward pass to a one-hot sample (matching eval-time - top-1 dispatch) while keeping the soft sample's gradient on backward. - """ - probs = self.gate(cond_cont, cond_cat) - if not (self.gumbel and self.training): - return probs - log_probs = torch.log(probs.clamp_min(1e-8)) - return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1) - - def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """(B,) hard expert index, used for eval-time grouped dispatch.""" - return self.gate(cond_cont, cond_cat).argmax(dim=-1) - - def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017).""" - importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) - return (importance.std() / (importance.mean() + 1e-8)) ** 2 - - def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: - """Optional supervised auxiliary loss shaping the router's own belief. - - Default: none (a scalar 0). Routers gating on an unobservable - pre-step quantity (e.g. ProcessRouter) override this. - """ - return torch.zeros((), device=cond_cont.device) - - def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """Optional auxiliary loss rewarding sharper (lower-entropy) routing.""" - norm_entropy, _ = self.gate_stats(cond_cont, cond_cat) - return norm_entropy - - def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for - the full explanation, unchanged in v0.3.0.""" - gate = self.gate(cond_cont, cond_cat) # (B, n_experts) - row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1) # (B,) - norm_entropy = row_entropy.mean() / math.log(self.n_experts) - importance = gate.sum(dim=0) # (n_experts,) - return norm_entropy, importance - - -ROUTER_REGISTRY: dict[str, type[Router]] = {} - - -def register_router(name: str): - def decorator(cls: type[Router]) -> type[Router]: - ROUTER_REGISTRY[name] = cls - return cls - - return decorator - - -def build_router(name: str, n_experts: int, **kwargs) -> Router: - """Factory: look up a `Router` subclass by name from the registry. - - Every registered router type is fed the same `router` config dict; - kwargs not declared by that type's constructor are silently dropped, so - per-type hyperparameters (e.g. EnergyRouter's `temperature`) can coexist - in one config without special-casing. - """ - if name not in ROUTER_REGISTRY: - raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}") - cls = ROUTER_REGISTRY[name] - accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"} - filtered = {k: v for k, v in kwargs.items() if k in accepted} - return cls(n_experts=n_experts, **filtered) - - -def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor: - """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient - bound used for EnergyRouter's `learn_width`/`learn_temperature` modes.""" - return lo + (hi - lo) * torch.sigmoid(raw) - - -def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float: - """Inverse of `_bounded_interp`, used once at construction to warm-start - `raw` so the initial effective width/temperature exactly equals `value`.""" - p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6) - return math.log(p / (1 - p)) - - -@register_router("energy") -class EnergyRouter(Router): - """Soft turn-on gate over normalized pre-step log-energy. - - Reads `cond_cont[:, energy_idx]` (ignores cond_cat). `gate(e) = - softmax_i(-(e - c_i)^2 / tau)`; as tau -> 0 this hardens to - nearest-center (Voronoi) selection, exactly what `top1` uses at eval. - """ - - def __init__( - self, - n_experts: int = 4, - temperature: float = 0.5, - learn_centers: bool = True, - energy_idx: int = 3, - centers_init: Sequence[float] | None = None, - learn_width: bool = False, - learn_temperature: bool = False, - width_min_ratio: float = 0.1, - width_max_ratio: float = 10.0, - ) -> None: - super().__init__(n_experts) - if learn_width and learn_temperature: - raise ValueError("learn_width and learn_temperature are mutually exclusive") - self.temperature = temperature - self.energy_idx = energy_idx - self.learn_width = learn_width - self.learn_temperature = learn_temperature - if learn_width or learn_temperature: - if not (width_min_ratio < 1.0 < width_max_ratio): - raise ValueError( - f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0" - ) - self._width_lo = width_min_ratio * temperature - self._width_hi = width_max_ratio * temperature - raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi) - if learn_width: - self.raw_width = nn.Parameter(torch.full((n_experts,), raw0)) - else: - self.raw_temperature = nn.Parameter(torch.tensor(raw0)) - if centers_init is None: - centers = torch.linspace(-2.0, 2.0, n_experts) - else: - if len(centers_init) != n_experts: - raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}") - centers = torch.tensor(list(centers_init), dtype=torch.float32) - if learn_centers: - self.centers = nn.Parameter(centers) - else: - self.register_buffer("centers", centers) - - def effective_width(self) -> torch.Tensor | float: - if self.learn_width: - return _bounded_interp(self.raw_width, self._width_lo, self._width_hi) - if self.learn_temperature: - return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi) - return self.temperature - - def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - e = cond_cont[:, self.energy_idx].unsqueeze(-1) # (B, 1) - d2 = (e - self.centers.unsqueeze(0)) ** 2 # (B, n_experts) - return torch.softmax(-d2 / self.effective_width(), dim=-1) - - -@register_router("pdg") -class PdgRouter(Router): - """Soft turn-on gate over a learned PDG embedding (own table, separate - from the trunk's `ConditionEncoder`). No supervision needed — PDG code - is already known at pre-step time.""" - - def __init__( - self, - n_experts: int, - pdg_vocab: int, - emb_dim: int = 8, - temperature: float = 0.5, - learn_centers: bool = True, - ) -> None: - super().__init__(n_experts) - self.temperature = temperature - self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) - centers = torch.randn(n_experts, emb_dim) * 0.1 - if learn_centers: - self.centers = nn.Parameter(centers) - else: - self.register_buffer("centers", centers) - - def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim) - d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1) # (B, n_experts) - return torch.softmax(-d2 / self.temperature, dim=-1) - - -@register_router("process") -class ProcessRouter(Router): - """Routes on the physics process expected to end the step — a post-step - outcome, so a small classifier over pre-step conditioning predicts it - (own pdg/material embeddings, separate from the trunk's ConditionEncoder). - `n_experts` doubles as the number of process classes. Supervised via - `classify_loss` against the true `process` label at train time only; - `gate`/`top1` never see it.""" - - def __init__( - self, - n_experts: int, - pdg_vocab: int, - mat_vocab: int, - emb_dim: int = 8, - hidden_dim: int = 64, - ) -> None: - super().__init__(n_experts) - self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) - self.mat_emb = nn.Embedding(mat_vocab, emb_dim) - self.classifier = nn.Sequential( - nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim), - nn.SiLU(), - nn.Linear(hidden_dim, n_experts), - ) - - def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - pdg_e = self.pdg_emb(cond_cat[:, 0]) - mat_e = self.mat_emb(cond_cat[:, 1]) - h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1) - return self.classifier(h) - - def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1) - - def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: - return F.cross_entropy(self.logits(cond_cont, cond_cat), labels) - - -class ComposedRouter(Router): - """Joint router over independent axes (e.g. energy x pdg), outer-product - gated. Not registered in `ROUTER_REGISTRY`; use `build_composed_router`.""" - - def __init__(self, routers: list[Router]) -> None: - if not routers: - raise ValueError("ComposedRouter needs at least one sub-router") - n_experts = 1 - for r in routers: - n_experts *= r.n_experts - super().__init__(n_experts) - self.routers = nn.ModuleList(routers) - - def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0) - for router in self.routers[1:]: - g = router.gate(cond_cont, cond_cat) # (B, n_i) - joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1) # (B, prod so far) - return joint - - def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: - total = torch.zeros((), device=cond_cont.device) - for router in self.routers: - total = total + router.classify_loss(cond_cont, cond_cat, labels) - return total - - -def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter: - """Build a `ComposedRouter` from a list of per-axis router specs — see - `_parse_composed_axes`.""" - routers = [ - build_router( - spec["type"], - spec["n_experts"], - **{ - **shared_kwargs, - **{k: v for k, v in spec.items() if k not in ("type", "n_experts")}, - }, - ) - for spec in specs - ] - return ComposedRouter(routers) - - -_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$") - - -def _parse_composed_axes(router_cfg: dict) -> list[dict]: - """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts. - - e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`, - `axis1_n_experts = 3`, `axis1_emb_dim = 8`. Axis indices must be - contiguous from 0. - """ - axes: dict[int, dict] = {} - for key, value in router_cfg.items(): - m = _AXIS_KEY_RE.match(key) - if m is None: - continue - idx, field = int(m.group(1)), m.group(2) - axes.setdefault(idx, {})[field] = value - missing = set(range(len(axes))) - axes.keys() - if missing: - raise ValueError(f"composed router config has gaps at axis indices {missing}") - return [axes[i] for i in range(len(axes))] - - -# Router types that read cond_cat's pdg index through their own -# nn.Embedding(pdg_vocab, ...), regardless of the trunk's particle -# conditioning mode — see _check_router_conditioning_compat. -_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process") - - -def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None: - """Reject a router axis that reintroduces a training-vocab PDG lookup - under `conditioning.particle.type = "physical"`. - - `PdgRouter`/`ProcessRouter` always build their own dataset-scoped - `nn.Embedding(pdg_vocab, ...)`, independent of `ConditionEncoder`'s - particle mode. Pairing either with `"physical"` would silently - reintroduce a training-menu-scoped lookup at the routing layer, - defeating the point of physical-property conditioning. Raised loudly at - model-build time. - """ - bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES)) - if bad and particle_conditioning == "physical": - raise ValueError( - f"router type(s) {bad} always use a training-vocab PDG embedding, " - "which is incompatible with conditioning.particle.type='physical' " - "(whose whole point is generalizing beyond that vocab) — pick a " - "different router type (e.g. 'energy') or use " - "conditioning.particle.type='embedding'." - ) - - -def _build_router_from_cfg( - router_cfg: dict, - pdg_vocab: int, - mat_vocab: int, - particle_conditioning: str = "embedding", -) -> Router: - """Resolve one stage's `router` config into a `Router`, single-axis or - composed. `gumbel` is set as a post-construction attribute (shared by - every router type, not a per-type constructor kwarg).""" - shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab) - if router_cfg["type"] == "composed": - axes = _parse_composed_axes(router_cfg) - _check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning) - router = build_composed_router(axes, **shared_vocab) - router.gumbel = bool(router_cfg.get("gumbel", False)) - return router - _check_router_conditioning_compat([router_cfg["type"]], particle_conditioning) - router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")} - router_kwargs.setdefault("pdg_vocab", pdg_vocab) - router_kwargs.setdefault("mat_vocab", mat_vocab) - router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs) - router.gumbel = bool(router_cfg.get("gumbel", False)) - return router - - -# --------------------------------------------------------------------------- -# Trunks -# --------------------------------------------------------------------------- - - -class ExpertTrunk(nn.Module): - """One small expert: `input_proj -> ResBlock stack -> out_proj`. - - Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR - tokens later (`noise_dim` in, `4 + type_dim` out), even though every - step-2/3 caller still has `in_dim == out_dim`. - """ - - def __init__( - self, - in_dim: int, - out_dim: int, - hidden_dim: int, - n_blocks: int, - cond_dim: int, - dropout: float = 0.0, - ) -> None: - super().__init__() - self.input_proj = nn.Linear(in_dim, hidden_dim) - self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)]) - self.out_proj = nn.Linear(hidden_dim, out_dim) - - def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: - x = self.input_proj(x) - for block in self.blocks: - x = block(x, cond) - return self.out_proj(x) - - -def _route_forward( - experts: nn.ModuleList, - router: Router, - x: torch.Tensor, - cond: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - training: bool, -) -> torch.Tensor: - """Shared dispatch for `RoutedTrunk`. - - Train mode: full mixture `sum_i weight_i * expert_i(x)` — always - N-expert dense compute, fully differentiable (`weight` is - `router.combine_weights`). Eval mode: grouped top-1 dispatch — each row - runs exactly one expert, the actual source of the per-call speedup. - """ - if training: - weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts) - out = torch.zeros(x.shape[0], experts[0].out_proj.out_features, device=x.device) - for i, expert in enumerate(experts): - out = out + weights[:, i : i + 1] * expert(x, cond) - return out - - idx = router.top1(cond_cont, cond_cat) # (B,) - out_dim = experts[0].out_proj.out_features - out = torch.zeros(x.shape[0], out_dim, device=x.device) - for i, expert in enumerate(experts): - mask = idx == i - if mask.any(): - out[mask] = expert(x[mask], cond[mask]) - return out - - -class Trunk(nn.Module): - """Interface implemented by `MonolithicTrunk`/`RoutedTrunk`: everything - downstream of the fused conditioning vector, i.e. the actual generative - trunk of a stage (`input_proj -> blocks -> out_proj`, monolithic or - expert-routed).""" - - def forward( - self, - x: torch.Tensor, - cond: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - ) -> torch.Tensor: - raise NotImplementedError - - -class MonolithicTrunk(Trunk): - def __init__( - self, - in_dim: int, - out_dim: int, - hidden_dim: int, - n_res_blocks: int, - cond_dim: int, - dropout: float = 0.0, - ) -> None: - super().__init__() - self.input_proj = nn.Linear(in_dim, hidden_dim) - self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_res_blocks)]) - self.out_proj = nn.Linear(hidden_dim, out_dim) - - def forward( - self, - x: torch.Tensor, - cond: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - ) -> torch.Tensor: - x = self.input_proj(x) - for block in self.blocks: - x = block(x, cond) - return self.out_proj(x) - - -class RoutedTrunk(Trunk): - def __init__( - self, - router: Router, - in_dim: int, - out_dim: int, - hidden_dim: int, - n_res_blocks: int, - cond_dim: int, - dropout: float = 0.0, - ) -> None: - super().__init__() - self.router = router - self.experts = nn.ModuleList( - [ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) for _ in range(router.n_experts)] - ) - - def forward( - self, - x: torch.Tensor, - cond: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - ) -> torch.Tensor: - return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training) - - -def build_trunk( - router: Router | None, - in_dim: int, - out_dim: int, - hidden_dim: int, - n_res_blocks: int, - cond_dim: int, - dropout: float = 0.0, -) -> Trunk: - if router is not None: - return RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) - return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) - - -# --------------------------------------------------------------------------- -# History encoders — stage-2 autoregressive only -# --------------------------------------------------------------------------- - - -class HistoryEncoder(nn.Module): - """Interface for stage-2 autoregressive per-token history summaries: - `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over - a full (teacher-forced) token sequence — used by training. `MarkovHistory` - and `AttentionHistory` are the two implementations. Inference - (`giant/sample.py`) generates one token at a - time and cannot afford `forward`'s per-step cost to be O(K) (attention - would then be O(K^2) over a rollout's k_max loop); encoders that need - incremental state for that path additionally implement `init_cache`/ - `step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since - its per-step cost is already O(1) (it only ever looks at the previous - token, not the full prefix).""" - - def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: - raise NotImplementedError - - -class MarkovHistory(HistoryEncoder): - """Summarizes the previous secondary's own `(energy_fraction, direction, - type_representation)` through one small MLP — the "markov" history: - token i+1 only ever sees token i plus the running scalars - (`remaining_frac`/`slot_idx`, fused in separately by - `Stage2Autoregressive._token_cond`), not the full prefix. - - At slot 0 (`has_prev` False) substitutes a learned start vector rather - than zeros — a reasonable default. - """ - - def __init__(self, in_dim: int, out_dim: int) -> None: - super().__init__() - self.start = nn.Parameter(torch.zeros(in_dim)) - self.mlp = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU()) - - def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: - start = self.start.view(1, 1, -1).expand_as(feat) - x = torch.where(has_prev.unsqueeze(-1), feat, start) - return self.mlp(x) - - -class _CausalAttnBlock(nn.Module): - """One pre-norm causal self-attention block for `AttentionHistory`. - - Exposes two forward paths that must agree (see - `test_attention_history_step_matches_forward` in `tests/test_network.py`): - `forward` — the full-sequence, causally-masked pass used for training; - `step` — an incremental pass for inference, given the *pre-attention* - normalized hidden states of every earlier position (`kv_cache`, i.e. - `norm1(x)` for positions `< t`, not `x` itself). Caching `norm1(x)` rather - than raw `x` is what makes `step` correct: this block's attention needs - exactly that quantity as keys/values, and `LayerNorm` has no cross-position - interaction, so recomputing it per position instead of caching it would - still be correct but pointlessly repeat work. The *next* block's cache is - built from a different sequence (this block's output), so each block owns - an independent cache entry. - """ - - def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None: - super().__init__() - self.norm1 = nn.LayerNorm(dim) - self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True) - self.norm2 = nn.LayerNorm(dim) - self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)) - - def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor: - h = self.norm1(x) - attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False) - x = x + attn_out - x = x + self.mlp(self.norm2(x)) - return x - - def step(self, x_new: torch.Tensor, kv_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]: - """`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None` - (first position) or `(B, T, dim)` — `norm1(x)` of every earlier - position at this same block. Returns `(out, new_kv_cache)`, `out` - being this position's block output (`(B, 1, dim)`, to feed the next - block's `step`), `new_kv_cache` the same cache extended by this - position (to reuse at this block's *next* `step` call).""" - h_new = self.norm1(x_new) - kv = h_new if kv_cache is None else torch.cat([kv_cache, h_new], dim=1) - attn_out, _ = self.attn(h_new, kv, kv, need_weights=False) - x = x_new + attn_out - x = x + self.mlp(self.norm2(x)) - return x, kv - - -class AttentionHistory(HistoryEncoder): - """Causal self-attention over the emitted-token prefix — the more - expressive alternative to `MarkovHistory`'s fixed previous-token-only - summary. `feat`/`has_prev` - follow the same shifted-by-one convention `MarkovHistory` and - `Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s - own `(energy_fraction, direction, type_representation)`, with a learned - start vector substituted at `has_prev == False` positions (only slot 0 in - practice — see `giant.training.stage2_inputs._ar_has_prev`). Causal masking then makes - position `i`'s output a function of `feat[:, 1:i+1]` — i.e. tokens - `0..i-1` — exactly the prefix available when predicting token `i`. - - `forward` is the parallel training path (one pass over the whole - teacher-forced sequence); `init_cache`/`step` are the incremental - inference path `giant/sample.py` uses, one new token per call, to avoid - re-encoding the whole prefix from scratch every slot — `step` must be - called exactly once per slot (its cache-extension is not idempotent), - so a slot's output must be reused for - every model call within that slot (`forward`'s ODE substeps, or a separate - `predict_type` call) rather than re-derived — see - `Stage2Autoregressive.history_step`. - """ - - def __init__(self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2) -> None: - super().__init__() - self.start = nn.Parameter(torch.zeros(in_dim)) - self.in_proj = nn.Linear(in_dim, out_dim) - self.blocks = nn.ModuleList([_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)]) - - def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: - start = self.start.view(1, 1, -1).expand_as(feat) - x = torch.where(has_prev.unsqueeze(-1), feat, start) - return self.in_proj(x) - - def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor: - B, K, _ = feat.shape - x = self._embed(feat, has_prev) - mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device) - for block in self.blocks: - x = block(x, mask) - return x - - def init_cache(self) -> list[torch.Tensor | None]: - return [None for _ in self.blocks] - - def step( - self, - token_feat: torch.Tensor, - has_prev: torch.Tensor, - cache: list[torch.Tensor | None], - ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: - """`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest - token's own features (what would be `feat[:, k]` in `forward`). - Advances every block's cache by this position and returns this - position's output (`(B, 1, out_dim)`, the correct history summary for - the NEXT slot) plus the updated cache.""" - x = self._embed(token_feat, has_prev) - new_cache: list[torch.Tensor | None] = [] - for block, kv in zip(self.blocks, cache): - x, kv_new = block.step(x, kv) - new_cache.append(kv_new) - return x, new_cache - - -# --------------------------------------------------------------------------- -# Stage models -# --------------------------------------------------------------------------- - - -def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int: - """Width of a single secondary slot's type slice — - `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else - `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are - `conditioning.particle.emb_dim` wide).""" - target = particle_type_cfg.get("target", "physical") - return PARTICLE_PHYS_DIM if target == "physical" else emb_dim - - -def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int: - """`Stage2OneShot`'s trunk output width. - - `target = "physical"` is untouched from v0.2/today: - `k_max * SEC_SLOT_DIM`, the type slice folded into the same - flow-matched/WGAN vector as the continuous stick/dir slots. - - `target` in `("onehot", "embedding")`: under `generator == "wgan"` the - type slice is still folded in (adversarial for onehot via ST-Gumbel, - already-continuous for embedding), just `emb_dim` wide instead of - `PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under - `generator in ("flow", "ddpm")` the type slice isn't part of this vector - at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk - only covers `k_max * CONT_SLOT_DIM`. - """ - target = particle_type_cfg.get("target", "physical") - if target == "physical": - return k_max * SEC_SLOT_DIM - if generator == "wgan": - return k_max * (CONT_SLOT_DIM + emb_dim) - return k_max * CONT_SLOT_DIM - - -class Stage1Model(nn.Module): - """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs - move it to stage 2, except for a migrated v0.2 checkpoint - (`n_sec_head_k_max` given), where it stays attached here - since that's where its weights live and what conditioning it was trained - against (see `_migrate_legacy_model_config`). - - `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — `conditioning.share_stages = true`: `build_models` - constructs one shared instance and passes it to both stages, halving the - conditioning parameter count and forcing a common representation.""" - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, - hidden_dim: int = 256, - n_res_blocks: int = 6, - cond_out_dim: int = 128, - x_dim: int = X_DIM, - dropout: float = 0.0, - generator: str = "flow", - time_dim: int = 64, - noise_dim: int = 64, - router: Router | None = None, - n_sec_head_k_max: int | None = None, - cond_enc: ConditionEncoder | None = None, - ) -> None: - super().__init__() - self.generator_kind = generator - self.noise_dim = noise_dim - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) - ) - has_time = generator in ("flow", "ddpm") - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - in_dim = noise_dim if generator == "wgan" else x_dim - self.trunk = build_trunk(router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout) - self.n_sec_head = None - if n_sec_head_k_max is not None: - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, n_sec_head_k_max + 1), - ) - - def forward( - self, - x_t: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - t: torch.Tensor | None = None, - ) -> torch.Tensor: - c_emb = self.cond_enc(cond_cont, cond_cat) - cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb - return self.trunk(x_t, cond, cond_cont, cond_cat) - - def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only - valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0 - configs predict n_sec from Stage2OneShot instead.""" - if self.n_sec_head is None: - raise RuntimeError( - "this Stage1Model has no n_sec_head — n_sec now lives on " - "stage 2 by default; this method only exists " - "for a migrated v0.2 checkpoint (n_sec.owner='stage1')" - ) - c_emb = self.cond_enc(cond_cont, cond_cat) - return self.n_sec_head(c_emb) - - -class Stage2OneShot(nn.Module): - """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour, - reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`, - step 4/5, not implemented yet). - - Owns `n_sec_head` by default unless `build_n_sec_head=False` - (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to - Stage1Model — see `_migrate_legacy_model_config`). - - `particle_type_cfg["target"]` (default `"physical"`) selects the - secondary-type mechanism: `"physical"` keeps the type slice folded into - the trunk's own - flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by - the caller via `stage2_trunk_sec_dim` — already reflects this). Under - `"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type - slice is predicted by a separate `type_head` instead (same shape pattern - as `n_sec_head`) — `sec_dim` then covers only the continuous - stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors. - Under `generator == "wgan"` the type slice stays folded into `sec_dim` - (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is - unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation. - - `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, - hidden_dim: int = 256, - n_res_blocks: int = 6, - cond_out_dim: int = 128, - context_dim: int = 64, - sec_dim: int = SEC_DIM, - x_dim: int = X_DIM, - dropout: float = 0.0, - generator: str = "wgan", - time_dim: int = 64, - noise_dim: int = 64, - k_max: int = K_MAX, - router: Router | None = None, - build_n_sec_head: bool = True, - particle_type_cfg: dict | None = None, - cond_enc: ConditionEncoder | None = None, - ) -> None: - super().__init__() - self.generator_kind = generator - self.noise_dim = noise_dim - self.k_max = k_max - self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) - self.type_dim = stage2_type_dim(self.particle_type_cfg, particle_cfg["emb_dim"]) - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) - ) - self.context_adapter = ContextAdapter(x_dim, context_dim) - self.fuse = nn.Sequential( - nn.Linear(cond_out_dim + context_dim, cond_out_dim), - nn.SiLU(), - ) - has_time = generator in ("flow", "ddpm") - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - in_dim = noise_dim if generator == "wgan" else sec_dim - self.trunk = build_trunk(router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout) - self.n_sec_head = None - if build_n_sec_head: - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, k_max + 1), - ) - self.type_head = None - target = self.particle_type_cfg.get("target", "physical") - if target != "physical" and generator in ("flow", "ddpm"): - emb_dim = particle_cfg["emb_dim"] - self.type_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, k_max * emb_dim), - ) - self._type_k_max = k_max - self._type_emb_dim = emb_dim - - def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: - base = self.cond_enc(cond_cont, cond_cat) - ctx = self.context_adapter(stage1_out) - return self.fuse(torch.cat([base, ctx], dim=-1)) - - def forward( - self, - x_t: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - t: torch.Tensor | None = None, - ) -> torch.Tensor: - c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) - cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb - return self.trunk(x_t, cond, cond_cont, cond_cat) - - def predict_n_sec( - self, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - ) -> torch.Tensor: - if self.n_sec_head is None: - raise RuntimeError( - "this Stage2OneShot has no n_sec_head — it belongs to a " - "migrated v0.2 checkpoint (n_sec.owner='stage1'); call " - "stage1.predict_n_sec(cond_cont, cond_cat) instead" - ) - c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) - return self.n_sec_head(c_emb) - - def predict_type( - self, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - ) -> torch.Tensor: - """`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or - vectors (`target="embedding"`) — only under `generator in ("flow", - "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s - own output instead (see class docstring).""" - if self.type_head is None: - raise RuntimeError( - "this Stage2OneShot has no type_head — either " - "particle_type.target='physical' (the type slice is part of " - "forward()'s own output) or generator='wgan' (the WGAN " - "trainer reads the type slice out of forward()'s output " - "directly instead)" - ) - c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out) - return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim) - - -class Stage2Autoregressive(nn.Module): - """Emits secondaries one at a time in descending-energy order, instead - of `Stage2OneShot`'s simultaneous - k_max-slot prediction. `history` selects `MarkovHistory` or - `AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only). - `teacher_forcing` handling lives entirely in the trainer - (`giant/train.py`), since it only affects how training inputs are - assembled, not this module's architecture. - - Under teacher forcing every token's conditioning is built from ground - truth, so a whole K-token sequence trains in one parallel batched pass: - `forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded - to `k_max`) — this also means a future one-token-at-a-time inference loop - (`K=1` per call, step 6) needs no interface change here. - - Two independent conditioning paths, mirroring `Stage2OneShot`'s - `_cond_embed` but split in two: `_base_cond` (`cond_enc` + - `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend - on token position; `_token_cond` additionally fuses in the history - encoding and two running scalars (remaining energy-budget fraction, - normalized slot index), and feeds `forward`/`predict_type`/the trunk. - - `cond_enc`, if given, is used in place of building a fresh - `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, - hidden_dim: int = 256, - n_res_blocks: int = 6, - cond_out_dim: int = 128, - context_dim: int = 64, - x_dim: int = X_DIM, - dropout: float = 0.0, - generator: str = "wgan", - time_dim: int = 64, - noise_dim: int = 64, - k_max: int = K_MAX, - router: Router | None = None, - build_n_sec_head: bool = True, - particle_type_cfg: dict | None = None, - history: str = "markov", - attn_n_heads: int = 4, - attn_n_layers: int = 2, - cond_enc: ConditionEncoder | None = None, - ) -> None: - super().__init__() - if history not in ("markov", "attention"): - raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'") - self.history_kind = history - self.generator_kind = generator - self.noise_dim = noise_dim - self.k_max = k_max - self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) - emb_dim = particle_cfg["emb_dim"] - self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim) - - self.cond_enc = ( - cond_enc - if cond_enc is not None - else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) - ) - self.context_adapter = ContextAdapter(x_dim, context_dim) - self.base_fuse = nn.Sequential( - nn.Linear(cond_out_dim + context_dim, cond_out_dim), - nn.SiLU(), - ) - - # Reuses conditioning.out_dim for the history encoder's own output - # width — there's no dedicated stage2_model.autoregressive key for - # this, a reasonable default rather than a design-doc-specified value. - history_dim = cond_out_dim - hist_in_dim = CONT_SLOT_DIM + self.type_dim - self.history_encoder: HistoryEncoder = ( - AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers) - if history == "attention" - else MarkovHistory(hist_in_dim, history_dim) - ) - token_fuse_in = cond_out_dim + context_dim + history_dim + 2 # +2: remaining_frac, slot_idx - self.token_fuse = nn.Sequential( - nn.Linear(token_fuse_in, cond_out_dim), - nn.SiLU(), - ) - - has_time = generator in ("flow", "ddpm") - self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None - merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim - token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim) - in_dim = noise_dim if generator == "wgan" else token_dim - self.trunk = build_trunk( - router, - in_dim, - token_dim, - hidden_dim, - n_res_blocks, - merged_cond_dim, - dropout, - ) - - self.n_sec_head = None - if build_n_sec_head: - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, k_max + 1), - ) - self.type_head = None - target = self.particle_type_cfg.get("target", "physical") - if target != "physical" and generator in ("flow", "ddpm"): - self.type_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, self.type_dim), - ) - - def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: - base = self.cond_enc(cond_cont, cond_cat) - ctx = self.context_adapter(stage1_out) - return self.base_fuse(torch.cat([base, ctx], dim=-1)) - - def _token_cond( - self, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - history_feat: torch.Tensor, - has_prev: torch.Tensor, - remaining_frac: torch.Tensor, - slot_idx: torch.Tensor, - hist: torch.Tensor | None = None, - ) -> torch.Tensor: - """`hist`, if given, overrides recomputing `self.history_encoder` - from `history_feat`/`has_prev` — the inference-time KV-cache path - (`Stage2Autoregressive.history_step`) precomputes it once per slot and - passes it in here so a slot's (possibly several) model calls — an ODE - loop's substeps, or a separate `predict_type` call — read the same - cached history instead of each re-deriving (and, under attention, - re-appending to the cache — see `AttentionHistory.step`'s docstring).""" - K = history_feat.size(1) - base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1) - ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1) - if hist is None: - hist = self.history_encoder(history_feat, has_prev) - scalars = torch.stack([remaining_frac, slot_idx], dim=-1) - return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1)) - - def init_history_cache(self): - """Inference-only incremental-decoding state for `self.history_encoder` - (`giant/sample.py`'s AR loop): `None` under `history="markov"` (its - per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or - `AttentionHistory.init_cache()` under `history="attention"`.""" - if isinstance(self.history_encoder, AttentionHistory): - return self.history_encoder.init_cache() - return None - - def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]: - """One inference slot's worth of history encoding: advances `cache` - (from `init_history_cache`, or a previous `history_step` call) by - `token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous - token, same convention `giant.sample.sample_secondaries_ar` already - threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is - this slot's history summary (pass it as `_token_cond`'s `hist=` to - every model call made for this slot), `new_cache` is what to pass into - the *next* slot's `history_step`. Must be called exactly once per - slot — see `AttentionHistory.step`'s docstring.""" - if isinstance(self.history_encoder, AttentionHistory): - return self.history_encoder.step(token_feat, has_prev, cache) - return self.history_encoder(token_feat, has_prev), cache - - def forward( - self, - x_t: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - history_feat: torch.Tensor, - has_prev: torch.Tensor, - remaining_frac: torch.Tensor, - slot_idx: torch.Tensor, - t: torch.Tensor | None = None, - hist: torch.Tensor | None = None, - ) -> torch.Tensor: - B, K = x_t.shape[0], x_t.shape[1] - c_emb = self._token_cond( - cond_cont, - cond_cat, - stage1_out, - history_feat, - has_prev, - remaining_frac, - slot_idx, - hist=hist, - ) - if self.time_emb is not None: - assert t is not None - t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1) - cond = torch.cat([t_emb, c_emb], dim=-1) - else: - cond = c_emb - x_flat = x_t.reshape(B * K, -1) - cond_flat = cond.reshape(B * K, -1) - cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1) - cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1) - out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat) - return out.view(B, K, -1) - - def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor: - if self.n_sec_head is None: - raise RuntimeError( - "this Stage2Autoregressive has no n_sec_head — it belongs to " - "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call " - "stage1.predict_n_sec(cond_cont, cond_cat) instead" - ) - return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out)) - - def predict_type( - self, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - history_feat: torch.Tensor, - has_prev: torch.Tensor, - remaining_frac: torch.Tensor, - slot_idx: torch.Tensor, - hist: torch.Tensor | None = None, - ) -> torch.Tensor: - if self.type_head is None: - raise RuntimeError( - "this Stage2Autoregressive has no type_head — either " - "particle_type.target='physical' (the type slice is part of " - "forward()'s own output) or generator='wgan' (the WGAN " - "trainer reads the type slice out of forward()'s output " - "directly instead)" - ) - c_emb = self._token_cond( - cond_cont, - cond_cat, - stage1_out, - history_feat, - has_prev, - remaining_frac, - slot_idx, - hist=hist, - ) - B, K, _ = c_emb.shape - return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim) - - -class CriticModel(nn.Module): - """Generator-agnostic WGAN-GP critic body: a scalar realism score, for - either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"` - mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as - `Stage2OneShot`). Used only when that stage's `generator == "wgan"`.""" - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - particle_cfg: dict, - material_cfg: dict, - in_dim: int, - hidden_dim: int = 256, - n_res_blocks: int = 6, - cond_out_dim: int = 128, - dropout: float = 0.0, - stage: str = "stage1", - context_dim: int = 64, - context_in_dim: int = X_DIM, - ) -> None: - super().__init__() - if stage not in ("stage1", "stage2"): - raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}") - self.stage = stage - self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim) - if stage == "stage2": - self.context_adapter = ContextAdapter(context_in_dim, context_dim) - self.fuse = nn.Sequential( - nn.Linear(cond_out_dim + context_dim, cond_out_dim), - nn.SiLU(), - ) - self.input_proj = nn.Linear(in_dim, hidden_dim) - self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)]) - self.out_norm = nn.LayerNorm(hidden_dim) - self.out_proj = nn.Linear(hidden_dim, 1) - - def forward( - self, - x: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor | None = None, - ) -> torch.Tensor: - base = self.cond_enc(cond_cont, cond_cat) - if self.stage == "stage2": - ctx = self.context_adapter(stage1_out) - cond = self.fuse(torch.cat([base, ctx], dim=-1)) - else: - cond = base - h = self.input_proj(x) - for block in self.blocks: - h = block(h, cond) - return self.out_proj(self.out_norm(h)).squeeze(-1) - - -# --------------------------------------------------------------------------- -# v0.2 -> v0.3 checkpoint migration -# --------------------------------------------------------------------------- - - -def _migrate_legacy_model_config(model_config: dict) -> dict: - """Translate a v0.2 checkpoint's flat `model_config` (giant/pipeline.py's - old shape: `hidden_dim`/`n_blocks`/`emb_dim`/`dropout`/`conditioning`/ - `router`/`mode`/... all at one level) into the nested - `{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model", - "stage2_model"}` shape `build_models` expects. - - Sets `stage2_model.n_sec.owner = "stage1"` so the n_sec_head weights a v0.2 - checkpoint carries on its Stage-1 module keep loading there instead of the new - default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's - own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not - just be labeled as such. - - Only the monolithic (non-routed) trunk shape is exercised by the step-2 - migration test; a routed v0.2 checkpoint still builds correctly here - (the router config passes through), but its - state dict isn't covered by `migrate_legacy_state_dict` below. - """ - m = model_config - conditioning_mode = m.get("conditioning", "embedding") - generator = m.get("mode", "flow") - hidden_dim = m.get("hidden_dim", 256) - n_blocks = m.get("n_blocks", 6) - emb_dim = m.get("emb_dim", EMB_DIM) - dropout = m.get("dropout", 0.1) - k_max = m.get("k_max", K_MAX) - noise_dim = m.get("noise_dim", 64) - router_cfg = dict(m.get("router") or {}) - reject_legacy_router_expert_sizing(router_cfg, source="this checkpoint's model_config.router") - router_cfg.setdefault("enabled", False) - - F = V02_FIXED_FACTS - cond_n_layers = F["conditioning.particle.n_layers"] # same fact for both axes - return { - "pdg_vocab": m["pdg_vocab"], - "mat_vocab": m["mat_vocab"], - "conditioning": { - "out_dim": F["conditioning.out_dim"], - "share_stages": False, - "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers}, - "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers}, - }, - "stage1_model": { - "active": F["stage1_model.active"], - "generator": generator, - "hidden_dim": hidden_dim, - "n_res_blocks": n_blocks, - "dropout": dropout, - "flow": {"time_dim": F["stage1_model.flow.time_dim"]}, - "ddpm": {"time_dim": F["stage1_model.ddpm.time_dim"]}, - "wgan": {"noise_dim": noise_dim}, - "router": dict(router_cfg), - }, - "stage2_model": { - "active": F["stage2_model.active"], - "decoder": F["stage2_model.decoder"], - "generator": generator, - "hidden_dim": hidden_dim, - "n_res_blocks": n_blocks, - "dropout": dropout, - "k_max": k_max, - "context_dim": F["stage2_model.context_dim"], - "n_sec": {"mode": "head", "owner": "stage1"}, - "particle_type": {"target": F["stage2_model.particle_type.target"]}, - "flow": {"time_dim": F["stage2_model.flow.time_dim"]}, - "ddpm": {"time_dim": F["stage2_model.ddpm.time_dim"]}, - "wgan": {"noise_dim": noise_dim}, - "router": {**router_cfg, "tie_to_stage1": False}, - }, - } - - -def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple[dict, dict]: - """Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`, - `SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new - `(Stage1Model, Stage2OneShot)` module structure produced by - `build_models(_migrate_legacy_model_config(model_config))`. - - Only the monolithic (non-routed) trunk shape is handled. - """ - - def _trunk_prefix(k: str) -> str: - if k.startswith(("input_proj.", "blocks.", "out_proj.")): - return f"trunk.{k}" - return k - - new_stage1 = {} - for k, v in old_stage1_sd.items(): - if k.startswith("n_sec_head."): - new_stage1[k] = v # stays top-level (n_sec.owner="stage1") - else: - new_stage1[_trunk_prefix(k)] = v - - new_stage2 = {} - for k, v in old_stage2_sd.items(): - if k.startswith("cond_enc.base."): - new_stage2["cond_enc." + k[len("cond_enc.base.") :]] = v - elif k.startswith("cond_enc.stage1_proj."): - new_stage2["context_adapter.proj." + k[len("cond_enc.stage1_proj.") :]] = v - elif k.startswith("cond_enc.fuse."): - new_stage2["fuse." + k[len("cond_enc.fuse.") :]] = v - else: - new_stage2[_trunk_prefix(k)] = v - - return new_stage1, new_stage2 - - -# --------------------------------------------------------------------------- -# Factories -# --------------------------------------------------------------------------- - - -def build_models(model_config: dict) -> dict[str, nn.Module | None]: - """Construct `{"stage1": ..., "stage2": ...}` from a config dict — either - the new nested shape (has a `"stage1_model"` key, plus `"pdg_vocab"`/ - `"mat_vocab"`/`"conditioning"` at the top level) or a v0.2 checkpoint's - flat `model_config`, auto-migrated via `_migrate_legacy_model_config`. - - A stage is `None` in the result when that stage's `active = False`. - `stage2_model.router.tie_to_stage1` shares stage 1's literal `Router` - instance rather than building a second, independently-parameterized one - (v0.2's actual — probably accidental — behaviour: two routers built from - one config with no semantic relationship between them). - - `conditioning.share_stages = true` builds one `ConditionEncoder` - instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/ - `Stage2Autoregressive`'s `cond_enc` param), instead of each stage - building its own — halving the conditioning parameter count and forcing a - common representation. `false` (default) keeps v0.2 behaviour: - independent instances with identical config but independent weights. - """ - cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config) - pdg_vocab = cfg["pdg_vocab"] - mat_vocab = cfg["mat_vocab"] - conditioning = cfg["conditioning"] - particle_cfg = conditioning["particle"] - material_cfg = conditioning["material"] - particle_conditioning = particle_cfg["type"] - 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_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 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 = 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 - n_sec_owner = s2_spec.n_sec.owner - n_sec_head_k_max = s2_spec.k_max if n_sec_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=s1_spec.hidden_dim, - n_res_blocks=s1_spec.n_res_blocks, - cond_out_dim=cond_out_dim, - dropout=s1_spec.dropout, - generator=generator, - 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 s2_spec.active: - decoder = s2_spec.decoder - router_cfg = cfg["stage2_model"].get("router") or {} - stage2_router: Router | None = 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 = 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 - n_sec_owner = s2_spec.n_sec.owner - k_max = s2_spec.k_max - particle_type_cfg = s2_spec.particle_type.to_dict() - - if decoder == "autoregressive": - 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=s2_spec.hidden_dim, - n_res_blocks=s2_spec.n_res_blocks, - cond_out_dim=cond_out_dim, - context_dim=s2_spec.context_dim, - dropout=s2_spec.dropout, - generator=generator, - time_dim=time_dim, - noise_dim=s2_spec.wgan.noise_dim, - k_max=k_max, - router=stage2_router, - build_n_sec_head=n_sec_owner != "stage1", - particle_type_cfg=particle_type_cfg, - 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: - sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, k_max, particle_cfg["emb_dim"]) - result["stage2"] = Stage2OneShot( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - particle_cfg=particle_cfg, - material_cfg=material_cfg, - hidden_dim=s2_spec.hidden_dim, - n_res_blocks=s2_spec.n_res_blocks, - cond_out_dim=cond_out_dim, - context_dim=s2_spec.context_dim, - sec_dim=sec_dim, - dropout=s2_spec.dropout, - generator=generator, - time_dim=time_dim, - noise_dim=s2_spec.wgan.noise_dim, - k_max=k_max, - router=stage2_router, - build_n_sec_head=n_sec_owner != "stage1", - particle_type_cfg=particle_type_cfg, - cond_enc=shared_cond_enc, - ) - - return result - - -def build_critics(model_config: dict) -> dict[str, nn.Module | None]: - """Construct `{"stage1": ..., "stage2": ...}` critics for `generator = - "wgan"` training. Training-only — never persisted for inference the way - `build_models`'s pair is. `None` for a stage that's inactive or not - WGAN.""" - cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config) - pdg_vocab = cfg["pdg_vocab"] - mat_vocab = cfg["mat_vocab"] - conditioning = cfg["conditioning"] - particle_cfg = conditioning["particle"] - material_cfg = conditioning["material"] - 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 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=s1_spec.hidden_dim, - n_res_blocks=s1_spec.n_res_blocks, - cond_out_dim=cond_out_dim, - dropout=s1_spec.dropout, - stage="stage1", - ) - - 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, - mat_vocab=mat_vocab, - particle_cfg=particle_cfg, - material_cfg=material_cfg, - in_dim=in_dim, - hidden_dim=s2_spec.hidden_dim, - n_res_blocks=s2_spec.n_res_blocks, - cond_out_dim=cond_out_dim, - dropout=s2_spec.dropout, - stage="stage2", - context_dim=s2_spec.context_dim, - ) - - return result +__all__ = [ + "AttentionHistory", + "ComposedRouter", + "ConditionEncoder", + "ContextAdapter", + "CriticModel", + "EnergyRouter", + "ExpertTrunk", + "HistoryEncoder", + "MarkovHistory", + "MonolithicTrunk", + "PdgRouter", + "ProcessRouter", + "ROUTER_REGISTRY", + "ResBlock", + "RoutedTrunk", + "Router", + "SinusoidalEmbedding", + "Stage1Model", + "Stage2Autoregressive", + "Stage2OneShot", + "Trunk", + "_CausalAttnBlock", + "_build_router_from_cfg", + "_check_router_conditioning_compat", + "_make_axis_mlp", + "_migrate_legacy_model_config", + "_parse_composed_axes", + "_route_forward", + "build_composed_router", + "build_critics", + "build_models", + "build_router", + "build_trunk", + "cat_col_layout", + "migrate_legacy_state_dict", + "register_router", + "stage2_trunk_sec_dim", + "stage2_type_dim", +] diff --git a/giant/model/routers.py b/giant/model/routers.py new file mode 100644 index 0000000..c5c69db --- /dev/null +++ b/giant/model/routers.py @@ -0,0 +1,375 @@ +"""Mixture-of-experts routing: `Router` base + registry, the four concrete +router types, and composed/config-driven construction — self-contained, no +dependency on any other `giant.model` submodule (issues.md Issue 8).""" + +import inspect +import math +import re +from collections.abc import Sequence + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from giant.constants import COND_DIM + +# --------------------------------------------------------------------------- +# Routers — carried over unchanged from v0.2 +# --------------------------------------------------------------------------- + + +class Router(nn.Module): + """Contract for a pluggable mixture-of-experts routing axis. + + Subclasses implement `gate` (soft partition-of-unity weights over + experts, used in train mode for a fully differentiable mixture); + `top1` and `balance_loss` have working defaults so a new routing axis + is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`. + """ + + def __init__(self, n_experts: int) -> None: + super().__init__() + self.n_experts = n_experts + self.gumbel = False + self.gumbel_tau = 1.0 + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B, n_experts) soft weights, rows summing to 1.""" + raise NotImplementedError + + def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B, n_experts) train-time expert-combination weights. + + Default (`gumbel=False`): identical to `gate()`. Opt-in + straight-through Gumbel-softmax (`gumbel=True`, train mode only): + hardens the forward pass to a one-hot sample (matching eval-time + top-1 dispatch) while keeping the soft sample's gradient on backward. + """ + probs = self.gate(cond_cont, cond_cat) + if not (self.gumbel and self.training): + return probs + log_probs = torch.log(probs.clamp_min(1e-8)) + return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1) + + def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """(B,) hard expert index, used for eval-time grouped dispatch.""" + return self.gate(cond_cont, cond_cat).argmax(dim=-1) + + def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017).""" + importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,) + return (importance.std() / (importance.mean() + 1e-8)) ** 2 + + def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """Optional supervised auxiliary loss shaping the router's own belief. + + Default: none (a scalar 0). Routers gating on an unobservable + pre-step quantity (e.g. ProcessRouter) override this. + """ + return torch.zeros((), device=cond_cont.device) + + def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + """Optional auxiliary loss rewarding sharper (lower-entropy) routing.""" + norm_entropy, _ = self.gate_stats(cond_cont, cond_cat) + return norm_entropy + + def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for + the full explanation, unchanged in v0.3.0.""" + gate = self.gate(cond_cont, cond_cat) # (B, n_experts) + row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1) # (B,) + norm_entropy = row_entropy.mean() / math.log(self.n_experts) + importance = gate.sum(dim=0) # (n_experts,) + return norm_entropy, importance + + +ROUTER_REGISTRY: dict[str, type[Router]] = {} + + +def register_router(name: str): + def decorator(cls: type[Router]) -> type[Router]: + ROUTER_REGISTRY[name] = cls + return cls + + return decorator + + +def build_router(name: str, n_experts: int, **kwargs) -> Router: + """Factory: look up a `Router` subclass by name from the registry. + + Every registered router type is fed the same `router` config dict; + kwargs not declared by that type's constructor are silently dropped, so + per-type hyperparameters (e.g. EnergyRouter's `temperature`) can coexist + in one config without special-casing. + """ + if name not in ROUTER_REGISTRY: + raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}") + cls = ROUTER_REGISTRY[name] + accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"} + filtered = {k: v for k, v in kwargs.items() if k in accepted} + return cls(n_experts=n_experts, **filtered) + + +def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor: + """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient + bound used for EnergyRouter's `learn_width`/`learn_temperature` modes.""" + return lo + (hi - lo) * torch.sigmoid(raw) + + +def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float: + """Inverse of `_bounded_interp`, used once at construction to warm-start + `raw` so the initial effective width/temperature exactly equals `value`.""" + p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6) + return math.log(p / (1 - p)) + + +@register_router("energy") +class EnergyRouter(Router): + """Soft turn-on gate over normalized pre-step log-energy. + + Reads `cond_cont[:, energy_idx]` (ignores cond_cat). `gate(e) = + softmax_i(-(e - c_i)^2 / tau)`; as tau -> 0 this hardens to + nearest-center (Voronoi) selection, exactly what `top1` uses at eval. + """ + + def __init__( + self, + n_experts: int = 4, + temperature: float = 0.5, + learn_centers: bool = True, + energy_idx: int = 3, + centers_init: Sequence[float] | None = None, + learn_width: bool = False, + learn_temperature: bool = False, + width_min_ratio: float = 0.1, + width_max_ratio: float = 10.0, + ) -> None: + super().__init__(n_experts) + if learn_width and learn_temperature: + raise ValueError("learn_width and learn_temperature are mutually exclusive") + self.temperature = temperature + self.energy_idx = energy_idx + self.learn_width = learn_width + self.learn_temperature = learn_temperature + if learn_width or learn_temperature: + if not (width_min_ratio < 1.0 < width_max_ratio): + raise ValueError( + f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0" + ) + self._width_lo = width_min_ratio * temperature + self._width_hi = width_max_ratio * temperature + raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi) + if learn_width: + self.raw_width = nn.Parameter(torch.full((n_experts,), raw0)) + else: + self.raw_temperature = nn.Parameter(torch.tensor(raw0)) + if centers_init is None: + centers = torch.linspace(-2.0, 2.0, n_experts) + else: + if len(centers_init) != n_experts: + raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}") + centers = torch.tensor(list(centers_init), dtype=torch.float32) + if learn_centers: + self.centers = nn.Parameter(centers) + else: + self.register_buffer("centers", centers) + + def effective_width(self) -> torch.Tensor | float: + if self.learn_width: + return _bounded_interp(self.raw_width, self._width_lo, self._width_hi) + if self.learn_temperature: + return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi) + return self.temperature + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + e = cond_cont[:, self.energy_idx].unsqueeze(-1) # (B, 1) + d2 = (e - self.centers.unsqueeze(0)) ** 2 # (B, n_experts) + return torch.softmax(-d2 / self.effective_width(), dim=-1) + + +@register_router("pdg") +class PdgRouter(Router): + """Soft turn-on gate over a learned PDG embedding (own table, separate + from the trunk's `ConditionEncoder`). No supervision needed — PDG code + is already known at pre-step time.""" + + def __init__( + self, + n_experts: int, + pdg_vocab: int, + emb_dim: int = 8, + temperature: float = 0.5, + learn_centers: bool = True, + ) -> None: + super().__init__(n_experts) + self.temperature = temperature + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + centers = torch.randn(n_experts, emb_dim) * 0.1 + if learn_centers: + self.centers = nn.Parameter(centers) + else: + self.register_buffer("centers", centers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim) + d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1) # (B, n_experts) + return torch.softmax(-d2 / self.temperature, dim=-1) + + +@register_router("process") +class ProcessRouter(Router): + """Routes on the physics process expected to end the step — a post-step + outcome, so a small classifier over pre-step conditioning predicts it + (own pdg/material embeddings, separate from the trunk's ConditionEncoder). + `n_experts` doubles as the number of process classes. Supervised via + `classify_loss` against the true `process` label at train time only; + `gate`/`top1` never see it.""" + + def __init__( + self, + n_experts: int, + pdg_vocab: int, + mat_vocab: int, + emb_dim: int = 8, + hidden_dim: int = 64, + ) -> None: + super().__init__(n_experts) + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + self.mat_emb = nn.Embedding(mat_vocab, emb_dim) + self.classifier = nn.Sequential( + nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim), + nn.SiLU(), + nn.Linear(hidden_dim, n_experts), + ) + + def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + pdg_e = self.pdg_emb(cond_cat[:, 0]) + mat_e = self.mat_emb(cond_cat[:, 1]) + h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1) + return self.classifier(h) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1) + + def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + return F.cross_entropy(self.logits(cond_cont, cond_cat), labels) + + +class ComposedRouter(Router): + """Joint router over independent axes (e.g. energy x pdg), outer-product + gated. Not registered in `ROUTER_REGISTRY`; use `build_composed_router`.""" + + def __init__(self, routers: list[Router]) -> None: + if not routers: + raise ValueError("ComposedRouter needs at least one sub-router") + n_experts = 1 + for r in routers: + n_experts *= r.n_experts + super().__init__(n_experts) + self.routers = nn.ModuleList(routers) + + def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0) + for router in self.routers[1:]: + g = router.gate(cond_cont, cond_cat) # (B, n_i) + joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1) # (B, prod so far) + return joint + + def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + total = torch.zeros((), device=cond_cont.device) + for router in self.routers: + total = total + router.classify_loss(cond_cont, cond_cat, labels) + return total + + +def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter: + """Build a `ComposedRouter` from a list of per-axis router specs — see + `_parse_composed_axes`.""" + routers = [ + build_router( + spec["type"], + spec["n_experts"], + **{ + **shared_kwargs, + **{k: v for k, v in spec.items() if k not in ("type", "n_experts")}, + }, + ) + for spec in specs + ] + return ComposedRouter(routers) + + +_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$") + + +def _parse_composed_axes(router_cfg: dict) -> list[dict]: + """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts. + + e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`, + `axis1_n_experts = 3`, `axis1_emb_dim = 8`. Axis indices must be + contiguous from 0. + """ + axes: dict[int, dict] = {} + for key, value in router_cfg.items(): + m = _AXIS_KEY_RE.match(key) + if m is None: + continue + idx, field = int(m.group(1)), m.group(2) + axes.setdefault(idx, {})[field] = value + missing = set(range(len(axes))) - axes.keys() + if missing: + raise ValueError(f"composed router config has gaps at axis indices {missing}") + return [axes[i] for i in range(len(axes))] + + +# Router types that read cond_cat's pdg index through their own +# nn.Embedding(pdg_vocab, ...), regardless of the trunk's particle +# conditioning mode — see _check_router_conditioning_compat. +_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process") + + +def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None: + """Reject a router axis that reintroduces a training-vocab PDG lookup + under `conditioning.particle.type = "physical"`. + + `PdgRouter`/`ProcessRouter` always build their own dataset-scoped + `nn.Embedding(pdg_vocab, ...)`, independent of `ConditionEncoder`'s + particle mode. Pairing either with `"physical"` would silently + reintroduce a training-menu-scoped lookup at the routing layer, + defeating the point of physical-property conditioning. Raised loudly at + model-build time. + """ + bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES)) + if bad and particle_conditioning == "physical": + raise ValueError( + f"router type(s) {bad} always use a training-vocab PDG embedding, " + "which is incompatible with conditioning.particle.type='physical' " + "(whose whole point is generalizing beyond that vocab) — pick a " + "different router type (e.g. 'energy') or use " + "conditioning.particle.type='embedding'." + ) + + +def _build_router_from_cfg( + router_cfg: dict, + pdg_vocab: int, + mat_vocab: int, + particle_conditioning: str = "embedding", +) -> Router: + """Resolve one stage's `router` config into a `Router`, single-axis or + composed. `gumbel` is set as a post-construction attribute (shared by + every router type, not a per-type constructor kwarg).""" + shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab) + if router_cfg["type"] == "composed": + axes = _parse_composed_axes(router_cfg) + _check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning) + router = build_composed_router(axes, **shared_vocab) + router.gumbel = bool(router_cfg.get("gumbel", False)) + return router + _check_router_conditioning_compat([router_cfg["type"]], particle_conditioning) + router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")} + router_kwargs.setdefault("pdg_vocab", pdg_vocab) + router_kwargs.setdefault("mat_vocab", mat_vocab) + router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs) + router.gumbel = bool(router_cfg.get("gumbel", False)) + return router diff --git a/giant/model/trunks.py b/giant/model/trunks.py new file mode 100644 index 0000000..2b9017a --- /dev/null +++ b/giant/model/trunks.py @@ -0,0 +1,155 @@ +"""Trunks: everything downstream of the fused conditioning vector — monolithic +or expert-routed (issues.md Issue 8).""" + +import torch +import torch.nn as nn + +from giant.model.layers import ResBlock +from giant.model.routers import Router + + +class ExpertTrunk(nn.Module): + """One small expert: `input_proj -> ResBlock stack -> out_proj`. + + Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR + tokens later (`noise_dim` in, `4 + type_dim` out), even though every + step-2/3 caller still has `in_dim == out_dim`. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + hidden_dim: int, + n_blocks: int, + cond_dim: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.input_proj = nn.Linear(in_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)]) + self.out_proj = nn.Linear(hidden_dim, out_dim) + + def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: + x = self.input_proj(x) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + +def _route_forward( + experts: nn.ModuleList, + router: Router, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + training: bool, +) -> torch.Tensor: + """Shared dispatch for `RoutedTrunk`. + + Train mode: full mixture `sum_i weight_i * expert_i(x)` — always + N-expert dense compute, fully differentiable (`weight` is + `router.combine_weights`). Eval mode: grouped top-1 dispatch — each row + runs exactly one expert, the actual source of the per-call speedup. + """ + if training: + weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts) + out = torch.zeros(x.shape[0], experts[0].out_proj.out_features, device=x.device) + for i, expert in enumerate(experts): + out = out + weights[:, i : i + 1] * expert(x, cond) + return out + + idx = router.top1(cond_cont, cond_cat) # (B,) + out_dim = experts[0].out_proj.out_features + out = torch.zeros(x.shape[0], out_dim, device=x.device) + for i, expert in enumerate(experts): + mask = idx == i + if mask.any(): + out[mask] = expert(x[mask], cond[mask]) + return out + + +class Trunk(nn.Module): + """Interface implemented by `MonolithicTrunk`/`RoutedTrunk`: everything + downstream of the fused conditioning vector, i.e. the actual generative + trunk of a stage (`input_proj -> blocks -> out_proj`, monolithic or + expert-routed).""" + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError + + +class MonolithicTrunk(Trunk): + def __init__( + self, + in_dim: int, + out_dim: int, + hidden_dim: int, + n_res_blocks: int, + cond_dim: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.input_proj = nn.Linear(in_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_res_blocks)]) + self.out_proj = nn.Linear(hidden_dim, out_dim) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + x = self.input_proj(x) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + +class RoutedTrunk(Trunk): + def __init__( + self, + router: Router, + in_dim: int, + out_dim: int, + hidden_dim: int, + n_res_blocks: int, + cond_dim: int, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.router = router + self.experts = nn.ModuleList( + [ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) for _ in range(router.n_experts)] + ) + + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training) + + +def build_trunk( + router: Router | None, + in_dim: int, + out_dim: int, + hidden_dim: int, + n_res_blocks: int, + cond_dim: int, + dropout: float = 0.0, +) -> Trunk: + if router is not None: + return RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) + return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) diff --git a/issues.md b/issues.md index c580f98..b4611db 100644 --- a/issues.md +++ b/issues.md @@ -47,7 +47,7 @@ architecture matrix grows, not about rot or breakage. | 5 | Inference bootstrap is duplicated verbatim between `predict` and `rollout` | **High** | Small | **Fixed** | | 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | **Fixed** | | 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | **Fixed** | -| 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | Open | +| 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | **Fixed** | | 9 | `scripts` is published as a top-level distribution package | Medium | Small | Open | | 10 | `torch.load(weights_only=False)` — checkpoints are arbitrary pickles | Low | Medium | Open | | 11 | Minor: `echo=print` threading, `particles → data.loader` layering | Low | Small | Open | @@ -1037,6 +1037,29 @@ protects against it, is the single riskiest version of this change. ## Issue 8 — `network.py` is 1745 lines holding three distinct modules +> **Status: Fixed.** `giant/model/network.py` (1742 lines) is now split exactly along the +> seams this issue identified: `giant/model/layers.py` (`SinusoidalEmbedding`, `ResBlock`, +> `ContextAdapter`, `_make_axis_mlp`), `encoders.py` (`cat_col_layout`, `ConditionEncoder`), +> `routers.py` (`Router` + registry + all 4 router types + composed/config parsing — +> self-contained, no dependency on any other new submodule), `trunks.py` (`Trunk`, +> `MonolithicTrunk`, `RoutedTrunk`, `ExpertTrunk`, `build_trunk`), `history.py` +> (`HistoryEncoder`, `MarkovHistory`, `AttentionHistory`, `_CausalAttnBlock` — +> self-contained), `models.py` (`stage2_type_dim`, `stage2_trunk_sec_dim`, `Stage1Model`, +> `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel`), `_legacy.py` +> (`_migrate_legacy_model_config`, `migrate_legacy_state_dict` — pure dict/state-dict +> translation, no dependency on the model classes), and `builders.py` (`build_models`, +> `build_critics`). `network.py` itself is now an 83-line re-export shim (`__all__` listing +> all 38 public + underscore-prefixed names any current caller reaches into), so none of +> the 22 external `from giant.model.network import ...` call sites (`sample.py`, +> `pipeline.py`, `checkpoint_io.py`, `training/trainers.py`, +> `analysis/router_gating.py`, and ~10 test files, several importing private names like +> `_build_router_from_cfg`/`_parse_composed_axes`) needed to change. Pure file move: no +> logic, signatures, or behavior changed — each new module carries only the imports it +> actually uses rather than a copy of the original file's full import block. `uv run +> pytest -q` (804 passed, unchanged), `uv run ruff check .`, `uv run ruff format --check +> .`, and `uv run ty check .` all clean. Everything below this point describes the +> pre-fix state and is kept for historical context. + **Severity: Medium. Effort: Small (mechanical).** **Location:** `giant/model/network.py`.