Merge pull request 'Make trunk architecture selectable via a registry (gitea #33)' (#48) from fix/issue-33 into master

Reviewed-on: #48

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")),
     (

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,

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,

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",

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)

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")

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