Make trunk architecture selectable via a registry (gitea #33) #48
Assignees
lars (Lars Bogner)
Clear assignees
lars (Lars Bogner)
No Assignees
Notifications
Due Date
No due date set.
Dependencies
No dependencies set.
Reference: lars/giant#48
Reference in New Issue
Block a user
Delete Branch "fix/issue-33"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
build_trunk hardcoded exactly two shapes (MonolithicTrunk/RoutedTrunk),
chosen only by whether a Router was built, with no way to select a
different trunk body architecture at all.
Deviates from the issue's literal proposal (a TRUNK_REGISTRY choosing
between "resmlp"/"moe" trunk shapes): during planning, decided that the
trunk body architecture and whether it's mixed are orthogonal, so the
registry (TRUNK_REGISTRY/register_trunk/build_expert_body in
giant/model/trunks.py) holds expert bodies only (today: "resmlp",
ExpertTrunk's existing input_proj -> ResBlock stack -> out_proj). Routing
stays exactly router.enabled/n_experts, untouched — a future transformer
body gets a mixture variant for free (trunk.type = "transformer" +
router.enabled = true) instead of needing a separate registry entry per
(body x routed/not) combination. MonolithicTrunk is deleted; the unrouted
case now returns the registry-selected body directly, preserving today's
exact state-dict keys (trunk.input_proj.* etc., not trunk.experts.0.*) —
required both for existing non-routed checkpoints and because
_legacy.py's migrate_legacy_state_dict already assumes that flat layout
for a v0.2 checkpoint.
New config leaf only: stage{1,2}_model.trunk.type: str = "resmlp"
(TrunkConfig). hidden_dim/n_res_blocks/dropout stay where they are today.
Nothing about router.enabled, config.migrate_config, _legacy.py, or the
CLI's --router flags changes — a v0.2-migrated config gets trunk.type =
"resmlp" automatically, reproducing current behaviour exactly. No CLI
flag added (matches the config.toml-only precedent set by
autoregressive.history/particle_type.target/n_sec.mode). No transformer
body and no "none"/"linear" body (gitea #45) in this change.
Full design rationale recorded on gitea #33 and #45 before implementation.
Co-Authored-By: Claude Opus 5 noreply@anthropic.com
build_trunk hardcoded exactly two shapes (MonolithicTrunk/RoutedTrunk), chosen only by whether a Router was built, with no way to select a different trunk body architecture at all. Deviates from the issue's literal proposal (a TRUNK_REGISTRY choosing between "resmlp"/"moe" trunk shapes): during planning, decided that the trunk *body* architecture and whether it's *mixed* are orthogonal, so the registry (TRUNK_REGISTRY/register_trunk/build_expert_body in giant/model/trunks.py) holds expert bodies only (today: "resmlp", ExpertTrunk's existing input_proj -> ResBlock stack -> out_proj). Routing stays exactly router.enabled/n_experts, untouched — a future transformer body gets a mixture variant for free (trunk.type = "transformer" + router.enabled = true) instead of needing a separate registry entry per (body x routed/not) combination. MonolithicTrunk is deleted; the unrouted case now returns the registry-selected body directly, preserving today's exact state-dict keys (trunk.input_proj.* etc., not trunk.experts.0.*) — required both for existing non-routed checkpoints and because _legacy.py's migrate_legacy_state_dict already assumes that flat layout for a v0.2 checkpoint. New config leaf only: stage{1,2}_model.trunk.type: str = "resmlp" (TrunkConfig). hidden_dim/n_res_blocks/dropout stay where they are today. Nothing about router.enabled, config.migrate_config, _legacy.py, or the CLI's --router flags changes — a v0.2-migrated config gets trunk.type = "resmlp" automatically, reproducing current behaviour exactly. No CLI flag added (matches the config.toml-only precedent set by autoregressive.history/particle_type.target/n_sec.mode). No transformer body and no "none"/"linear" body (gitea #45) in this change. Full design rationale recorded on gitea #33 and #45 before implementation. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>