Trunk architecture is not selectable (hardcoded ResMLP/MoE choice) #33

Open
opened 2026-08-13 15:07:03 +02:00 by lars · 0 comments
Owner

build_trunk (giant/model/trunks.py:144-155) hardcodes exactly two shapes,
both of them Linear → N × ResBlock → Linear, chosen by
router is None. The config exposes only hidden_dim, n_res_blocks,
dropout. There is no way to try a transformer trunk, a gated/GLU MLP, a
width-varying stack, or a deliberately trivial linear trunk as an ablation
baseline, without editing models.py.

Proposal: TRUNK_REGISTRY + register_trunk, and a
[stage1_model.trunk] / [stage2_model.trunk] block keyed on type:

[stage1_model.trunk]
type = "resmlp"        # "resmlp" | "moe" | "transformer" | "linear"
hidden_dim = 256
n_blocks = 6

router.enabled = true becomes trunk.type = "moe" with the router block
nested under it, which is the honest shape — a routed trunk is a kind of
trunk
, not a boolean modifier on the default one. (migrate_config maps the
old spelling forward; the deep-merge special-case that keeps router.enabled
from dropping the rest of the router defaults then goes away, since type keys
don't have that problem.)

Payoff: the single change that most directly delivers "the network should
be very customizable". Everything else in this batch of modularity issues is
plumbing by comparison.


Migrated from issues.md (v0.3.0 branch review, 2026-08-13), Issue 8.

`build_trunk` (`giant/model/trunks.py:144-155`) hardcodes exactly two shapes, both of them `Linear → N × ResBlock → Linear`, chosen by `router is None`. The config exposes only `hidden_dim`, `n_res_blocks`, `dropout`. There is no way to try a transformer trunk, a gated/GLU MLP, a width-varying stack, or a deliberately trivial `linear` trunk as an ablation baseline, without editing `models.py`. **Proposal:** `TRUNK_REGISTRY` + `register_trunk`, and a `[stage1_model.trunk]` / `[stage2_model.trunk]` block keyed on `type`: ```toml [stage1_model.trunk] type = "resmlp" # "resmlp" | "moe" | "transformer" | "linear" hidden_dim = 256 n_blocks = 6 ``` `router.enabled = true` becomes `trunk.type = "moe"` with the router block nested under it, which is the honest shape — a routed trunk *is a kind of trunk*, not a boolean modifier on the default one. (`migrate_config` maps the old spelling forward; the deep-merge special-case that keeps `router.enabled` from dropping the rest of the router defaults then goes away, since `type` keys don't have that problem.) **Payoff:** the single change that most directly delivers "the network should be very customizable". Everything else in this batch of modularity issues is plumbing by comparison. --- Migrated from `issues.md` (v0.3.0 branch review, 2026-08-13), Issue 8.
lars added the architecturemodularity labels 2026-08-13 15:07:03 +02:00
lars added a new dependency 2026-08-13 15:23:29 +02:00
Sign in to join this conversation.