diff --git a/giant/config.py b/giant/config.py index 52246a6..71c83de 100644 --- a/giant/config.py +++ b/giant/config.py @@ -336,6 +336,29 @@ class RouterConfig: } +@dataclass(frozen=True) +class TrunkConfig: + """`stage1_model.trunk`/`stage2_model.trunk`: selects the trunk's expert + *body* architecture from `giant.model.trunks.TRUNK_REGISTRY` (default + `"resmlp"` — today's only body, `input_proj -> ResBlock stack -> + out_proj`). Orthogonal to whether that body is mixed: mixing is still + controlled entirely by `router.enabled`/`router.n_experts` on the same + stage, unaffected by this block. A future body's own hyperparameters + (e.g. a transformer's `n_heads`/`n_layers`) would get their own sibling + field here, matching how `flow`/`ddpm`/`wgan` already coexist selected by + `generator`.""" + + type: str = "resmlp" + + @classmethod + def from_dict(cls, d: dict | None) -> "TrunkConfig": + d = d or {} + return cls(type=d.get("type", "resmlp")) + + def to_dict(self) -> dict: + return {"type": self.type} + + @dataclass(frozen=True) class Stage2RouterConfig(RouterConfig): # true: stage 2 shares stage 1's Router module instance, so expert i in @@ -509,6 +532,7 @@ class Stage1ModelConfig: ddpm: DdpmConfig = field(default_factory=DdpmConfig) wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig) router: RouterConfig = field(default_factory=RouterConfig) + trunk: TrunkConfig = field(default_factory=TrunkConfig) @classmethod def from_dict(cls, d: dict | None) -> "Stage1ModelConfig": @@ -524,6 +548,7 @@ class Stage1ModelConfig: ddpm=DdpmConfig.from_dict(d.get("ddpm")), wgan=Stage1WganConfig.from_dict(d.get("wgan")), router=RouterConfig.from_dict(d.get("router")), + trunk=TrunkConfig.from_dict(d.get("trunk")), ) def to_dict(self) -> dict: @@ -538,6 +563,7 @@ class Stage1ModelConfig: "ddpm": self.ddpm.to_dict(), "wgan": self.wgan.to_dict(), "router": self.router.to_dict(), + "trunk": self.trunk.to_dict(), } @@ -575,6 +601,7 @@ class Stage2ModelConfig: ddpm: DdpmConfig = field(default_factory=DdpmConfig) wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig) router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig) + trunk: TrunkConfig = field(default_factory=TrunkConfig) @classmethod def from_dict(cls, d: dict | None) -> "Stage2ModelConfig": @@ -597,6 +624,7 @@ class Stage2ModelConfig: ddpm=DdpmConfig.from_dict(d.get("ddpm")), wgan=Stage2WganConfig.from_dict(d.get("wgan")), router=Stage2RouterConfig.from_dict(d.get("router")), + trunk=TrunkConfig.from_dict(d.get("trunk")), ) def to_dict(self) -> dict: @@ -618,6 +646,7 @@ class Stage2ModelConfig: "ddpm": self.ddpm.to_dict(), "wgan": self.wgan.to_dict(), "router": self.router.to_dict(), + "trunk": self.trunk.to_dict(), } @@ -1400,6 +1429,8 @@ _OUT_DIR_NAME_CANDIDATES = [ "particle_type_target", _path_candidate("stage2_model.particle_type.target", "pt-"), ), + ("stage1_trunk_type", _path_candidate("stage1_model.trunk.type", "s1t-")), + ("stage2_trunk_type", _path_candidate("stage2_model.trunk.type", "s2t-")), ("stage1_router", _router_candidate("stage1_model", "s1")), ("stage2_router", _router_candidate("stage2_model", "s2")), ( diff --git a/giant/model/builders.py b/giant/model/builders.py index 553b376..f5c689d 100644 --- a/giant/model/builders.py +++ b/giant/model/builders.py @@ -83,6 +83,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: time_dim=time_dim, noise_dim=s1_spec.wgan.noise_dim, router=stage1_router, + trunk_type=s1_spec.trunk.type, n_sec_head_k_max=n_sec_head_k_max, cond_enc=shared_cond_enc, ) @@ -121,6 +122,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: noise_dim=s2_spec.wgan.noise_dim, k_max=k_max, router=stage2_router, + trunk_type=s2_spec.trunk.type, build_n_sec_head=n_sec_owner != "stage1", particle_type_cfg=particle_type_cfg, history=ar_cfg.history, @@ -148,6 +150,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: noise_dim=s2_spec.wgan.noise_dim, k_max=k_max, router=stage2_router, + trunk_type=s2_spec.trunk.type, build_n_sec_head=n_sec_owner != "stage1", particle_type_cfg=particle_type_cfg, cond_enc=shared_cond_enc, diff --git a/giant/model/models.py b/giant/model/models.py index ea00db1..ee3def7 100644 --- a/giant/model/models.py +++ b/giant/model/models.py @@ -92,6 +92,7 @@ class Stage1Model(nn.Module): time_dim: int = 64, noise_dim: int = 64, router: Router | None = None, + trunk_type: str = "resmlp", n_sec_head_k_max: int | None = None, cond_enc: ConditionEncoder | None = None, ) -> None: @@ -107,7 +108,7 @@ class Stage1Model(nn.Module): 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.trunk = build_trunk(router, trunk_type, 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( @@ -185,6 +186,7 @@ class Stage2OneShot(nn.Module): noise_dim: int = 64, k_max: int = K_MAX, router: Router | None = None, + trunk_type: str = "resmlp", build_n_sec_head: bool = True, particle_type_cfg: dict | None = None, cond_enc: ConditionEncoder | None = None, @@ -211,7 +213,9 @@ class Stage2OneShot(nn.Module): 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.trunk = build_trunk( + router, trunk_type, 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( @@ -328,6 +332,7 @@ class Stage2Autoregressive(nn.Module): noise_dim: int = 64, k_max: int = K_MAX, router: Router | None = None, + trunk_type: str = "resmlp", build_n_sec_head: bool = True, particle_type_cfg: dict | None = None, history: str = "markov", @@ -380,6 +385,7 @@ class Stage2Autoregressive(nn.Module): in_dim = noise_dim if generator == "wgan" else token_dim self.trunk = build_trunk( router, + trunk_type, in_dim, token_dim, hidden_dim, diff --git a/giant/model/network.py b/giant/model/network.py index c71f976..2cb2341 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -36,12 +36,14 @@ from giant.model.routers import ( register_router, ) from giant.model.trunks import ( + TRUNK_REGISTRY, ExpertTrunk, - MonolithicTrunk, RoutedTrunk, Trunk, _route_forward, + build_expert_body, build_trunk, + register_trunk, ) __all__ = [ @@ -54,7 +56,6 @@ __all__ = [ "ExpertTrunk", "HistoryEncoder", "MarkovHistory", - "MonolithicTrunk", "PdgRouter", "ProcessRouter", "ROUTER_REGISTRY", @@ -65,6 +66,7 @@ __all__ = [ "Stage1Model", "Stage2Autoregressive", "Stage2OneShot", + "TRUNK_REGISTRY", "Trunk", "_CausalAttnBlock", "_build_router_from_cfg", @@ -75,12 +77,14 @@ __all__ = [ "_route_forward", "build_composed_router", "build_critics", + "build_expert_body", "build_models", "build_router", "build_trunk", "cat_col_layout", "migrate_legacy_state_dict", "register_router", + "register_trunk", "resolve_type_n_classes", "stage2_trunk_sec_dim", "stage2_type_dim", diff --git a/giant/model/trunks.py b/giant/model/trunks.py index 2b9017a..ac31942 100644 --- a/giant/model/trunks.py +++ b/giant/model/trunks.py @@ -1,5 +1,13 @@ -"""Trunks: everything downstream of the fused conditioning vector — monolithic -or expert-routed (issues.md Issue 8).""" +"""Trunks: everything downstream of the fused conditioning vector — a +registrable expert *body* architecture (`TRUNK_REGISTRY`/`register_trunk`), +used standalone or mixed by a `Router` (issues.md Issue 8; trunk-selectability +gitea #33). + +Whether a body is mixed is orthogonal to which body it is: `RoutedTrunk` +builds `router.n_experts` instances of whichever body `trunk_type` names, so +a future body (e.g. a transformer) automatically gets a mixture variant for +free — no separate "routed transformer trunk" class needed. +""" import torch import torch.nn as nn @@ -7,9 +15,42 @@ import torch.nn as nn from giant.model.layers import ResBlock from giant.model.routers import Router +TRUNK_REGISTRY: dict[str, type[nn.Module]] = {} + +def register_trunk(name: str): + def decorator(cls: type[nn.Module]) -> type[nn.Module]: + TRUNK_REGISTRY[name] = cls + return cls + + return decorator + + +def build_expert_body( + name: str, + in_dim: int, + out_dim: int, + hidden_dim: int, + n_blocks: int, + cond_dim: int, + dropout: float = 0.0, +) -> nn.Module: + """Factory: look up a registered trunk body by name and construct one + instance of it — used both for a standalone (unrouted) trunk and for each + expert inside a `RoutedTrunk`.""" + if name not in TRUNK_REGISTRY: + raise ValueError(f"unknown trunk type {name!r}; available: {sorted(TRUNK_REGISTRY)}") + cls = TRUNK_REGISTRY[name] + return cls(in_dim, out_dim, hidden_dim, n_blocks, cond_dim, dropout) + + +@register_trunk("resmlp") class ExpertTrunk(nn.Module): - """One small expert: `input_proj -> ResBlock stack -> out_proj`. + """`input_proj -> ResBlock stack -> out_proj` — the registered `"resmlp"` + trunk body. Used both standalone (no router: `forward`'s `cond_cont`/ + `cond_cat` are accepted and ignored, satisfying the `Trunk` interface + directly with no wrapper class) and as one expert inside a `RoutedTrunk` + (`_route_forward` calls it with just `(x, cond)`). 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 @@ -26,11 +67,18 @@ class ExpertTrunk(nn.Module): dropout: float = 0.0, ) -> None: super().__init__() + self.out_dim = out_dim 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: + def forward( + self, + x: torch.Tensor, + cond: torch.Tensor, + cond_cont: torch.Tensor | None = None, + cond_cat: torch.Tensor | None = None, + ) -> torch.Tensor: x = self.input_proj(x) for block in self.blocks: x = block(x, cond) @@ -55,13 +103,13 @@ def _route_forward( """ 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) + out = torch.zeros(x.shape[0], experts[0].out_dim, 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_dim = experts[0].out_dim out = torch.zeros(x.shape[0], out_dim, device=x.device) for i, expert in enumerate(experts): mask = idx == i @@ -71,10 +119,10 @@ def _route_forward( 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).""" + """Interface implemented by a standalone trunk body (any `TRUNK_REGISTRY` + entry, e.g. `ExpertTrunk`) and by `RoutedTrunk`: everything downstream of + the fused conditioning vector, i.e. the actual generative trunk of a + stage.""" def forward( self, @@ -86,38 +134,11 @@ class Trunk(nn.Module): 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, + trunk_type: str, in_dim: int, out_dim: int, hidden_dim: int, @@ -128,7 +149,10 @@ class RoutedTrunk(Trunk): 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)] + [ + build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) + for _ in range(router.n_experts) + ] ) def forward( @@ -143,13 +167,23 @@ class RoutedTrunk(Trunk): def build_trunk( router: Router | None, + trunk_type: str, in_dim: int, out_dim: int, hidden_dim: int, n_res_blocks: int, cond_dim: int, dropout: float = 0.0, -) -> Trunk: +) -> nn.Module: + """Build a stage's trunk: `trunk_type` (a `TRUNK_REGISTRY` key, e.g. + `"resmlp"`) selects the expert body architecture; `router`, if given, + wraps `router.n_experts` instances of that body in a `RoutedTrunk` + mixture — otherwise a single body is returned directly (no wrapper + class), which is what makes an unrouted trunk's state-dict keys land + directly under `trunk.*` instead of `trunk.experts.0.*` (see + `giant.model._legacy.migrate_legacy_state_dict`, which assumes exactly + this flat layout for a v0.2 monolithic checkpoint). + """ 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) + return RoutedTrunk(router, trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) + return build_expert_body(trunk_type, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) diff --git a/tests/test_config.py b/tests/test_config.py index 9fdce45..e51a555 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -51,6 +51,7 @@ def test_giant_config_to_dict_matches_default_config(): gconfig.Stage2WganConfig, gconfig.RouterConfig, gconfig.Stage2RouterConfig, + gconfig.TrunkConfig, gconfig.NSecConfig, gconfig.ParticleTypeConfig, gconfig.AutoregressiveConfig, @@ -75,6 +76,15 @@ def test_stage2_model_config_defaults_match_documented_v030_intent(): assert spec.particle_type.target == "onehot" +def test_trunk_config_defaults_to_resmlp_for_both_stages(): + """gitea #33: a v0.2-migrated / pre-existing config with no `trunk` key + at all must reproduce today's behaviour exactly.""" + assert gconfig.Stage1ModelConfig().trunk.type == "resmlp" + assert gconfig.Stage2ModelConfig().trunk.type == "resmlp" + assert gconfig.DEFAULT_CONFIG["stage1_model"]["trunk"]["type"] == "resmlp" + assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["type"] == "resmlp" + + def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips(): """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim" — the default must stay 0 so an existing config.toml with no @@ -911,6 +921,21 @@ def test_validate_config_keys_skips_meta_section(): gconfig.validate_config_keys(cfg) # must not raise +def test_validate_config_keys_allows_trunk_type(): + cfg = _cfg_with(**{"stage1_model.trunk.type": "resmlp", "stage2_model.trunk.type": "resmlp"}) + gconfig.validate_config_keys(cfg) # must not raise + + +def test_validate_config_keys_rejects_unknown_trunk_key(): + cfg = _cfg_with(**{"stage1_model.trunk.type_o": "resmlp"}) # typo for type + try: + gconfig.validate_config_keys(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "stage1_model.trunk.type_o" in str(e) + assert "type" in str(e) + + def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path): path = tmp_path / "config.toml" path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n") diff --git a/tests/test_router.py b/tests/test_router.py index d733218..3c77c64 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -5,9 +5,10 @@ import torch from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM from giant.model.network import ( + TRUNK_REGISTRY, ComposedRouter, EnergyRouter, - MonolithicTrunk, + ExpertTrunk, PdgRouter, ProcessRouter, ROUTER_REGISTRY, @@ -15,6 +16,7 @@ from giant.model.network import ( Stage1Model, Stage2OneShot, build_composed_router, + build_expert_body, build_models, build_router, ) @@ -144,6 +146,22 @@ def test_build_router_unknown_type_raises(): raise AssertionError("expected ValueError for unknown router type") +# ── TRUNK_REGISTRY / build_expert_body ────────────────────────────────────── + + +def test_trunk_registry_has_resmlp(): + assert "resmlp" in TRUNK_REGISTRY + assert TRUNK_REGISTRY["resmlp"] is ExpertTrunk + + +def test_build_expert_body_unknown_type_raises(): + try: + build_expert_body("nonexistent", in_dim=4, out_dim=4, hidden_dim=8, n_blocks=1, cond_dim=4) + except ValueError: + return + raise AssertionError("expected ValueError for unknown trunk type") + + # ── EnergyRouter learn_width / learn_temperature ──────────────────────────── @@ -993,8 +1011,10 @@ def test_build_models_monolith_when_router_absent(): stage1, stage2 = models["stage1"], models["stage2"] assert isinstance(stage1, Stage1Model) assert isinstance(stage2, Stage2OneShot) - assert isinstance(stage1.trunk, MonolithicTrunk) - assert isinstance(stage2.trunk, MonolithicTrunk) + assert not isinstance(stage1.trunk, RoutedTrunk) + assert not isinstance(stage2.trunk, RoutedTrunk) + assert isinstance(stage1.trunk, ExpertTrunk) + assert isinstance(stage2.trunk, ExpertTrunk) def test_build_models_monolith_when_router_disabled(): @@ -1006,8 +1026,10 @@ def test_build_models_monolith_when_router_disabled(): models = build_models(cfg) stage1, stage2 = models["stage1"], models["stage2"] assert stage1 is not None and stage2 is not None - assert isinstance(stage1.trunk, MonolithicTrunk) - assert isinstance(stage2.trunk, MonolithicTrunk) + assert not isinstance(stage1.trunk, RoutedTrunk) + assert not isinstance(stage2.trunk, RoutedTrunk) + assert isinstance(stage1.trunk, ExpertTrunk) + assert isinstance(stage2.trunk, ExpertTrunk) def test_build_models_routed_when_enabled(): @@ -1040,6 +1062,21 @@ def test_build_models_routed_when_enabled(): assert len(stage2.trunk.experts) == 4 +def test_build_models_explicit_resmlp_trunk_type_matches_default(): + """stage1_model.trunk.type = 'resmlp' is the default's spelled-out + equivalent, not a behaviour change — gitea #33.""" + default_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2) + explicit_cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2, trunk={"type": "resmlp"}) + default_stage1 = build_models(default_cfg)["stage1"] + explicit_stage1 = build_models(explicit_cfg)["stage1"] + assert default_stage1 is not None and explicit_stage1 is not None + assert type(default_stage1.trunk) is type(explicit_stage1.trunk) is ExpertTrunk + assert default_stage1.trunk.input_proj.weight.shape == explicit_stage1.trunk.input_proj.weight.shape + default_params = sum(p.numel() for p in default_stage1.parameters()) + explicit_params = sum(p.numel() for p in explicit_stage1.parameters()) + assert default_params == explicit_params + + def test_build_models_routed_pair_is_drop_in_for_sample_flow(): """Exercise the exact calling convention giant/sample.py uses.""" from giant.sample import sample_flow, sample_secondaries