10 Commits

Author	SHA1	Message	Date
lars	4b2e0ba98e	Give CriticModel a registry-built trunk and StageModel base (gitea #57) ... CriticModel was the one stage-shaped class left out of the trunk-registry (gitea #33), block-conditioning-registry (gitea #34), and StageModel-base (gitea #39) refactors: it hand-rolled a plain ResBlock stack, so a routed/FiLM/AdaLN trunk was available to every generative stage model except the critic competing against them under WGAN-GP. CriticModel now subclasses StageModel (reusing its cond_enc construction, and a stage-2 context-fusion helper factored out of Stage2OneShot onto the base) and builds its body via build_trunk (output width 1) instead of a bespoke ResBlock loop, so trunk.type/trunk.block_conditioning now affect the critic too. Each stage's critic inherits its own generator's trunk config rather than a new critic_trunk config key, mirroring the existing critic_hidden_dim/critic_n_res_blocks "0 = inherit from generator" pattern. Router mixing (MoE) for the critic stays out of scope. Since CriticModel is training-only and never persisted for inference, and WGAN-GP is still unbenchmarked, its state_dict shape has no back-compat burden. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-24 09:48:26 +02:00
lars	c1c4957e2f	Implement n_sec.mode = "stop_token" for the AR secondary decoder (gitea #40) ... Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way: a one-shot n_sec_head classifier over conditioning alone, run before any secondary token existed, with the AR loop then always executing k_max slots and discarding the tail. This adds a real per-slot EOS mechanism instead: - Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head (n_sec.mode = "stop_token" builds no n_sec_head at all). - sample_secondaries_ar accepts n_sec_pred=None to drive generation off the stop head instead of a pre-resolved count: each row stops the first slot its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the default, threshold at 0 — or "sample", a Bernoulli draw), and the whole batch loop breaks once every row has stopped, so cost scales with the realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly (the scheduled-sampling self-sample path) is unchanged. - resolve_n_sec returns None for a stop-token decoder instead of raising; rollout.py/cli.py/validate.py now derive the realized count from sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling, rather than resolving it up front — a no-op reordering under every other n_sec.mode, where sec_valid was already built from n_sec_pred. - Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds the per-slot target/mask (one slot wider than the existing token-content sec_mask, since the stop slot itself needs supervision) and StageTrainer._stop_loss trains it with masked BCE, gated on stop_head exactly like _n_sec_loss gates on n_sec_head. Wired into both the flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop head replaces n_sec_head under this mode, so no new weight key. - validate_config now accepts stop_token (requires decoder="autoregressive" and n_sec.owner="stage2") instead of always rejecting it. Decisions made during planning: stop_sampling defaults to "greedy" for deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding new config keys, since the two heads never coexist. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-17 10:44:14 +02:00
lars	c4b12b5e7a	Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts (gitea #38) ... build_models/build_critics parsed model_config into frozen dataclasses (ConditioningConfig, Stage2ModelConfig, ...) but then threw the parsed sub-objects away and passed the original raw dicts (conditioning["particle"], s2_spec.particle_type.to_dict()) down into ConditionEncoder/StageModel/etc, which re-read them with their own hardcoded .get(key, default) fallbacks — each an independent copy of a fact the dataclass already stated once. Worst instance: giant/training/trainers.py:236 converted an already-parsed ParticleTypeConfig back into a dict for no reason. Threads ConditioningAxisConfig (particle_cfg/material_cfg) and ParticleTypeConfig (particle_type_cfg) as the actual dataclass instances through every signature that used to type them dict: ConditionEncoder, StageModel/CriticModel, resolve_type_n_classes/stage2_type_dim/ stage2_trunk_sec_dim, giant/model/builders.py, giant/sample.py, giant/training/stage2_inputs.py, giant/training/trainers.py (StageSpec/ StageTrainer), giant/pipeline.py, giant/rollout.py, giant/validate.py — so ty now catches a misspelled field instead of it silently falling back. No config-schema change: config.toml/checkpoint model_config keep the same nested-dict shape; only what happens after the existing X.from_dict(...) parse changes. User-confirmed scope decision: both axes (particle_cfg/material_cfg and particle_type_cfg), not just the more heavily-duplicated particle_type_cfg axis, and not stopping at the two most literal parse-then-discard round trips — matching the issue's own proposal. Preserved-default decision: StageModel's particle_type_cfg=None sentinel (hit only by direct/test construction — build_models always passes an explicit particle_type) still resolves to ParticleTypeConfig(target= "physical"), not ParticleTypeConfig()'s own target="onehot" config-file default — switching it would have silently grown an unused, gradient-less type_head on every test that constructs Stage2OneShot/Stage2Autoregressive without particle_type_cfg=, breaking their "every param has a grad" checks. New tests in tests/test_network.py: ConditionEncoder/StageModel store the exact ConditioningAxisConfig/ParticleTypeConfig instance passed in (identity, not just equality) — no internal dict round-trip — and build_models's output carries real dataclass instances end to end, not the plain dicts it produced before this fix. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 15:03:55 +02:00
lars	593c5f4d34	Deduplicate n_sec_head/type_head MLPs into build_mlp_head (gitea #36) ... The same two-layer classifier head (Linear(cond_out_dim, hidden_dim // 2) -> SiLU -> Linear(hidden_dim // 2, out_dim)) was hand-rolled five times in giant/model/models.py: Stage1Model.n_sec_head, Stage2OneShot.n_sec_head/ .type_head, and Stage2Autoregressive.n_sec_head/.type_head. The `// 2` ratio and fixed 2-layer depth were undocumented magic numbers, and both n_sec accuracy and secondary-species accuracy are known weak spots that were untunable independently of the trunk they hang off. Adds `build_mlp_head(in_dim, out_dim, hidden, depth, act)` to giant/model/layers.py (depth=1 is a bare Linear; depth>=2 matches the old hardcoded shape exactly), and a new `HeadConfig` (hidden_ratio, depth) dataclass in giant/config.py, wired in as `stage1_model.heads.n_sec` and `stage2_model.heads.{n_sec,type}` — split per head type (not one shared block per stage) since n_sec and species prediction are called out as separate weak spots that may want independent capacity. Defaults (hidden_ratio=0.5, depth=2) reproduce the old hardcoded architecture bit-for-bit, so every existing config.toml and migrated v0.2 checkpoint is unaffected; no changes were needed to migrate_config or the legacy migration surfaces. No new CLI flags, matching how other nested sub-config (router.*, trunk.*) is set via config.toml rather than per-field flags. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 10:57:47 +02:00
lars	f8722e347e	Add an Objective registry for the flow/ddpm/wgan generator choice (gitea #32) ... generator ∈ {"flow", "ddpm", "wgan"} was tested as a bare string in ~45 sites across models.py, sample.py, builders.py, trainers.py, and stage2_inputs.py, each independently re-deriving one of five consequences of the choice (needs a time embedding? what does the trunk take as input? is the type slice folded into the trunk output? which sampler? which loss?). giant/model/objectives.py adds an Objective ABC + OBJECTIVE_REGISTRY + build_objective factory, mirroring routers.py's Router pattern, and every bare-string site now goes through it (needs_time, is_adversarial, folds_type_slice, trunk_in_dim, build_schedule, stage1_loss/stage2_loss). Per discussion: FlowDDPMStageTrainer and WGANStageTrainer stay separate classes rather than merging into one StageTrainer as the issue's sketch proposed — their training loops are genuinely different shapes (single loss vs. dual G/D step with gradient penalty/n_critic/ST-Gumbel), and trainers.py is the least-covered-by-fast-tests part of the codebase, so a full merge was judged out of proportion to this issue's risk budget. FlowDDPMStageTrainer's own loss dispatch (flow vs ddpm, one-shot vs AR) does move onto the objective, so a future non-adversarial objective (rectified flow, consistency distillation) is still a one-file, zero-trainer-edits addition. No config-schema change — stage{1,2}_model.generator stays the persisted string, just looked up in the registry instead of string-compared. An unrecognized generator value now fails fast with a clear ValueError instead of silently falling through some bare-string checks and not others (same behavior build_router/build_trunk already have for their own type keys). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 10:10:58 +02:00
lars	0f95e0eaae	Make ResBlock's conditioning-injection mechanism selectable (gitea #34) ... ResBlock injected conditioning exactly one way — h = linear1(h) + cond_proj(cond), a conditional bias, the weakest standard option for a model whose entire job is to be conditional. Adds BLOCK_REGISTRY (giant/model/layers.py), mirroring the TRUNK_REGISTRY/ROUTER_REGISTRY registry+factory idiom (gitea #33), with two new drop-in alternatives: FilmResBlock (per-channel scale+shift modulating the norm output, zero-init so conditioning has no effect at construction) and AdaLNResBlock (DiT-style AdaLN-Zero — the norm's own affine is replaced by a conditioning-derived scale/shift, plus a zero-init gate on the residual branch, making the block the exact identity function at init). Selected per stage via a new stage{1,2}_model.trunk.block_conditioning config leaf ("add" | "film" | "adaln", default "add"), threaded through build_trunk/build_expert_body/RoutedTrunk and the three stage model constructors. Default stays "add" and ResBlock's body is unchanged, so existing configs/checkpoints are bit-identical to before this change. Decided during planning: the new field lives on the existing TrunkConfig rather than a new top-level block/blocks config section; the WGAN CriticModel (which builds its own ResBlock stack outside TRUNK_REGISTRY) and the issue's mentioned blocks.norm/blocks.activation axes are both left out of scope. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-14 09:41:50 +02:00
lars	f3f7645bf7	Make trunk architecture selectable via a registry (gitea #33) ... 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>	2026-08-14 09:16:55 +02:00
lars	32aa5a5f92	Decouple secondary-species vocabulary from conditioning.particle.emb_dim (gitea #29) ... conditioning.particle.emb_dim and stage2_model.particle_type.target="onehot"'s class count were silently the same number everywhere (pipeline.py's PDG top-N map build, Stage2OneShot/Stage2Autoregressive's type head, StageSpec's training loss width, the checkpoint's shared pdg_topn_map), fixing the secondary-species vocabulary at whatever width the unrelated physical-conditioning MLP happened to use — the exact vocabulary the v0.3.0 pivot exists to fix. Adds stage2_model.particle_type.n_classes (default 0 = inherit conditioning.particle.emb_dim, preserving today's behavior and every existing checkpoint) and a single resolve_type_n_classes helper used everywhere the coupling used to be implicit. Splits the checkpoint's shared pdg_topn_map into a conditioning-only pdg_topn_map and a new sec_type_topn_map, built independently through the existing (axis, n_classes)-keyed setup cache (no extra scan when they still resolve to the same N) and threaded through giant predict/giant rollout's decode path. A checkpoint with no sec_type_topn_map key (pre-#29) falls back to reusing pdg_topn_map, reproducing the old shared behavior exactly. Decided with the user during planning: commit directly on this branch; represent the split as an additive sec_type_topn_map checkpoint key rather than conditionally reusing pdg_topn_map; build the two top-N maps independently rather than the issue's proposed build-at-max-and-slice, since the setup cache already avoids redundant scans across runs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-13 16:11:14 +02:00
lars	da717971b6	Honour wgan.critic_hidden_dim/critic_n_res_blocks in build_critics (gitea #28) ... build_critics always sized a WGAN critic off the generator's own hidden_dim/n_res_blocks, silently discarding the documented 0=inherit sentinel on stage{1,2}_model.wgan.critic_hidden_dim/critic_n_res_blocks (the same convention critic_lr already honoured). Now both keys are read with the 0 -> inherit fallback, and stage-scoped-only CLI flags (--stage{1,2}-critic-hidden-dim/--stage{1,2}-critic-n-res-blocks) are added -- no shared alias, since critic sizing is an architectural per-stage knob like --hidden-dim/--n-res-blocks, not a shared training hyperparameter like --n-critic/--gp-weight/--noise-dim/--critic-lr. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-13 15:36:26 +02:00
lars	72f5a891bf	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.	2026-08-13 10:21:13 +02:00