23 Commits

Author	SHA1	Message	Date
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	899ca3a7d5	Validate stage2_model.autoregressive.order in validate_config (gitea #30) ... order was documented as single-valued ("energy_desc" only, placeholder for a future alternative ordering) but validate_config only checked its siblings history/teacher_forcing, so e.g. order = "energy_asc" was silently accepted and trained as if it were energy_desc. Add the missing check alongside the other two, gated the same way (only meaningful under stage2_model.decoder = "autoregressive"). Also updates the stale reason string on the pre-existing _KNOWN_UNUSED allow-list entry for this key in tests/test_config_consumed_keys.py, since half of it ("validate_config ... never [checks] order") is no longer true after this fix — the key stays allow-listed because validate_config itself isn't in that test's build/train/rollout consumer whitelist. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-13 15:40:47 +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	c3fc768b40	Reject stage2_model.stage1_context = 'sampled' as unimplemented (issues.md Issue 1) ... trainers.py unconditionally trains stage 2 against the ground-truth stage-1 output (stage1_ctx = x1_s1.detach()), but 'sampled' was accepted by validate_config, stored in config.toml and the checkpoint's model_config, and silently trained identically to 'truth' — mislabeling every downstream artifact for a run launched with --stage2-stage1-context sampled. Mirrors the existing stop_token validate_config pattern. User chose the immediate fix (reject loudly) over the proper fix (actually implement sampled context), which is scoped to Issue 16. Also updates the _KNOWN_UNUSED reason for stage2_model.stage1_context (added by Issue 5's consumed-keys audit) to reflect that the value is now rejected rather than silently accepted, and drops the now-invalid --stage2-stage1-context sampled case from test_stage2_only_knobs (a full CLI invocation) — that flag's plumbing is still covered at the overrides-dict level by test_overrides_from_flags_stage2_only_knobs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-13 14:57:21 +02:00
lars	e6261cea03	Unify the two v0.2->v0.3 migration surfaces (issues.md Issue 6) ... giant/config.py:migrate_config (config.toml) and giant/model/network.py:_migrate_legacy_model_config (checkpoint model_config) independently hand-maintained the same v0.2 facts and an identical router expert-sizing rejection. Extract the shared knowledge into a new leaf module, giant/_migration.py (V02_MODEL_KEY_TO_STAGES, V02_FIXED_FACTS, reject_legacy_router_expert_sizing), consumed by both. Also replace NSecConfig's legacy-only, nullable legacy_owner sentinel (living in an extra: dict catch-all) with a normal, always-set owner: str = "stage2" field, so build_models reads one concrete two-valued key instead of branching on a legacy marker. Record in CLAUDE.md that v0.2 checkpoint-loading support has no expiry decided yet, since /ceph still holds pre-v0.3.0 checkpoints. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-13 09:56:22 +02:00
lars	2bfb1ab056	Extract giant train/new-run's CLI override mapping into a table-driven function (issues.md Issues 3 & 4) ... train()'s ~140-line hand-written flag->config translation (three different ad hoc "more specific flag wins" patterns) and new_run()'s near-verbatim copy are replaced by a shared FlagSpec/FLAG_SPECS table and overrides_from_flags() in config.py, reused by both commands. This makes the override/precedence logic directly unit-testable without CliRunner, closing coverage gaps that had zero tests (e.g. --emb-dim/--conditioning dual-axis fan-out, three of four WGAN knob legs, --stage2-generator overriding --mode, router's stage1-only asymmetry). No CLI flags, help text, or precedence semantics changed -- `giant train --help`/`giant new-run --help` are byte-identical before and after, and all previously-passing CliRunner tests still pass unmodified. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-12 15:04:49 +02:00
lars	01acbfed61	Add unknown-key validation to config.toml merge (issues.md Issue 2) ... A typo like `n_res_block` for `n_res_blocks` previously merged cleanly, passed validate_config, and silently trained a model that didn't match config.toml's documented settings. merge_cli_overrides now rejects any key not present in DEFAULT_CONFIG's schema via validate_config_keys, with a did-you-mean suggestion, while still allowing the genuinely dynamic composed-router axis keys and centers_init. Checkpoint model_config loading is untouched, so old checkpoints keep loading. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-12 14:47:29 +02:00
lars	9bf5874308	Make config dataclasses the single source of truth for DEFAULT_CONFIG ... DEFAULT_CONFIG and build_models/build_critics/StageSpec.from_config's inline .get(key, default) fallbacks had already drifted: two keys (stage2_model.decoder, stage2_model.particle_type.target) resolved differently depending on whether a config dict came from merge_cli_overrides (fully populated, correct) or was hand-built and partial (fell back to stale v0.2-shaped literals). Introduce frozen dataclasses (GiantConfig and its nested blocks) in giant/config.py as the actual single declaration of every default; DEFAULT_CONFIG is now generated from them instead of hand-maintained, and build_models, build_critics, and StageSpec.from_config consume the dataclasses instead of duplicating literal fallbacks, so this class of drift can't recur. Router/n_sec sub-blocks keep an `extra` catch-all for their genuinely dynamic keys (composed-router axes, runtime-seeded centers_init, legacy_owner). Fixing the fallback surfaced the same latent bug in two existing partial-config callers that had been silently depending on it: a test fixture in test_train.py and scripts/warm_setup_cache.py's minimal cfg (now merged against DEFAULT_CONFIG instead of hand-rolled, closing the gap for good). See issues.md Issue 1. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-12 14:35:14 +02:00
lars	55332db67a	Bump ruff line-length to 120 and reformat ... Rejoins lines that only wrapped because they exceeded the old 88-char limit; ruff check and the full test suite (725 passed) are unaffected.	2026-08-12 13:33:09 +02:00
lars	878e9ddca3	Delete docs/v0.3.0-design.md and strip all references to it ... The design doc and its followups doc are no longer needed as a live reference now that the v0.3.0 redesign is implemented — comments and docstrings across the codebase cited it extensively (file path, "design doc §X.Y", "decision N", or bare "§X.Y" section numbers) as design rationale. Removed docs/ and edited every citing comment/docstring to drop the now-dangling reference while keeping the substantive explanation next to it. CLAUDE.md's v0.3.0 roadmap bullet loses its trailing pointer to the deleted file. Verified: no remaining "docs/v0.3.0", "design doc", "decision N", or "§N.N" references (repo-wide grep); ruff and ty clean; full test suite on the heaviest-touched modules (network, sample, rollout, migration, config, train) passes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-10 11:19:02 +02:00
lars	da7cde3ef9	v0.3.0 post-implementation audit: resolve all 9 tracked discrepancies ... Works through docs/v0.3.0-followups.md item by item, closing the gap between the design doc and the shipped v0.3.0-stage2-autoregressive code: 1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2 dispatch, stage-2 particle-type-class marginal. 2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run. 3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/ train instead of the hardcoded K_MAX constant. 4. Mixed conditioning.particle.type / conditioning.material.type support end-to-end (data pipeline + dwarf warm-cache). 5. conditioning.share_stages = true: one shared ConditionEncoder instance across both stages. 6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2 (was silently unimplemented). 7. giant predict/rollout: implement conditioning.*.type = "onehot" via the checkpoint's saved pdg_topn_map/mat_topn_map. 8. network.py's checkpoint-path model_config migration now fails loudly on non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's TOML-load path (§4.2). 9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a rollout-capable checkpoint (§9). Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics), mostly a test-helper dict-unpack pattern that made every unrelated constructor keyword look like a type error. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-07 16:12:58 +02:00
lars	200c6d243b	v0.3.0 step 7: AttentionHistory (KV-cached) + scheduled/never teacher forcing ... AttentionHistory (giant/model/network.py) adds causal self-attention over the emitted-secondary prefix as the alternative to MarkovHistory, with a parallel forward() for training and an init_cache()/step() KV-cache path for sample.py's per-slot AR inference loop, wired into Stage2Autoregressive via history="attention". giant/train.py adds _stage2_tf_prob and _assemble_stage2_ar_inputs_scheduled, mixing ground-truth history with a detached sample_secondaries_ar self-sample per slot so teacher_forcing="scheduled"/"never" close the train/inference gap teacher_forcing="always" always avoided; wired into both stage-2 AR trainers. config.py's validate_config no longer rejects these two previously unimplemented schema values. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-07 13:15:56 +02:00
lars	c9d255b1c5	v0.3.0 step 5: Stage2Autoregressive (history=markov) + §11.4 grad instrumentation ... Replaces the Stage2Autoregressive stub with a real per-token secondary decoder: MarkovHistory summarizes the previous secondary, remaining-energy fraction and slot index round out the per-token conditioning, and the existing Trunk/MonolithicTrunk/RoutedTrunk machinery is reused unchanged by batching all K_MAX tokens together under teacher forcing (one parallel pass, no new trunk code). build_models/build_critics wire it in; the WGAN critic stays whole-sequence, so build_critics needs no AR-specific path. train.py's FlowDDPMStageTrainer/WGANStageTrainer gain a decoder branch, sharing optimizer/EMA/checkpoint machinery with the one-shot path. _assemble_stage2_real is now defined in terms of the new unflattened _assemble_stage2_ar_target helper, removing a near-duplicate branch. Also lands the §11.4 differentiability validation-obligation instrumentation (trunk-gradient norm from the particle-type slice vs. the continuous slices, for generator=wgan + particle_type.target=onehot) via backward hooks in _relax_onehot_type_slice, decoder-agnostic and surfaced as two new metrics.csv columns. This also fixes the standing regression where any config not explicitly overriding decoder="one_shot" crashed at build_models, since stage2_model.decoder defaults to "autoregressive" — confirmed by removing tests/test_pipeline.py's now-stale override so the default config runs end-to-end against real synthetic data. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-07 09:36:49 +02:00
lars	9112e845e0	v0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite ... Replaces train.py's single global training loop with a StageTrainer hierarchy (FlowDDPMStageTrainer, WGANStageTrainer) — one per active stage, each owning its own optimizer/LR schedule/EMA and reading only the shared batch tuple (stage 2 always teacher-forces on the ground-truth x1_s1, so stages never need each other's output at train time). Supports every stage1/stage2 generator combination, including the design doc's headline mixed case (stage1=flow + stage2=wgan) and its reverse, plus stage1-only/stage2-only ablation runs, routed+gumbel stages, and checkpoint save/resume. metrics.csv/wandb logging are stage-prefixed. validate_marginals calls are guarded with a one-time warning and a Wasserstein-magnitude fallback for wgan best-checkpoint selection, since giant/sample.py still assumes stage1 always owns n_sec_head (decision 1 moved it to stage 2 by default) — deferred to design doc step 6, not silently papered over. pipeline.py's run_setup_stage/run_train_job now read the new nested config directly; the dangling resolve_expert_dims call and the --mode wgan --router rejection are both gone (routed WGAN works). cli.py's train/new-run build correctly-shaped config overrides (architecture flags -> stage1_model only per the approved decision; --mode/--n-critic/--gp-weight/--critic-lr broadcast to both stages, matching migrate_config's own precedent and avoiding a regression on the common --mode case); predict/rollout's dangling build_models tuple-unpack is fixed; new-run now tags config_version, fixing a bug where a re-loaded v0.3 config.toml would have been silently corrupted by migrate_config mistaking it for v0.2. config.py's validate_config rejects mixed particle/material conditioning types for now (ConditionEncoder supports it, the data pipeline in giant/data/transforms.py doesn't yet). analysis/render.py and router_gating.py handle both the new nested model_config shape and legacy flat checkpoints. scripts/warm_setup_cache.py updated for run_setup_stage's new signature. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-06 11:31:49 +02:00
lars	eb6dd27406	v0.3.0 step 1: new nested config schema, v0.2 migration shim ... Replace the single global train.mode + [model] block with the four top-level blocks docs/v0.3.0-design.md specifies ([conditioning], [stage1_model], [stage2_model], [train]), so Stage 1 and Stage 2 can run independent generative objectives and Stage 2 can train standalone. - migrate_config translates old config.toml/checkpoint dicts on load, so nothing on /ceph goes dead; loudly rejects non-zero expert_hidden_dim/expert_n_blocks, which v0.3.0 no longer supports. - merge_cli_overrides/save_config generalize from one hardcoded nesting level (model.router) to arbitrary recursive depth. - default_out_dir_name candidates move to dotted paths against the new schema, with per-stage router/generator discriminators. - validate_config adds cross-block checks the per-block schema can't express (particle_type.target=embedding needs a matching conditioning mode, tie_to_stage1 needs an active stage 1, etc). - resolve_expert_dims is deleted (experts always inherit the stage's hidden_dim/n_res_blocks now) — pipeline.py/cli.py callers are left dangling on purpose, to be updated in the network.py/train.py steps that follow. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-06 10:15:20 +02:00
lars	5eec4c250a	Add gumbel/learn_centers/learn_width/learn_temperature to out-dir naming ... Extends default_out_dir_name's non-default-field convention to the router's new gumbel combine-weight flag and its learnable-knob toggles, so gumbel sweep configs (learn_centers on/off, learn_width, learn_temperature) resolve to distinguishable checkpoint directory names instead of colliding on the same r-<type><n> token. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-07-30 16:42:56 +02:00
lars	b51eafcfa5	Add opt-in straight-through Gumbel-softmax combine weights to MoE router ... Trains the routed trunk's forward combination as a hard one-hot sample (matching eval-time top-1 dispatch exactly) while keeping a smooth gradient on the backward pass, targeting the train/eval mismatch identified as a likely contributor to experts overlapping instead of partitioning in the first energy-router rollout benchmark. Off by default (model.router.gumbel); existing routed configs/checkpoints are unaffected. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-07-30 16:29:20 +02:00
lars	1115451c8e	Make default checkpoint out_dir name reflect only non-default hyperparams ... Previously the same fixed 7 fields (mode/hidden_dim/n_blocks/emb_dim/ conditioning/lr/batch_size) were always baked into the name, even for a vanilla run, and router config wasn't represented at all. Now default_out_dir_name only includes fields that differ from DEFAULT_CONFIG, adds router/seed/epochs as candidates, and caps at 6 shown fields with a hashed overflow suffix for heavily-swept configs.	2026-07-29 11:21:02 +02:00
lars	539b6f61e1	Add test coverage for resolve_expert_dims ... Covers the default-config 0-sentinel inheritance path (the exact bug fixed by 969c5c6, previously untested since every router test in test_router.py passes expert_hidden_dim/expert_n_blocks explicitly), plus missing-key inheritance, full override, and partial override. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-07-29 10:42:44 +02:00
lars	8cebc4809d	Apply ruff format and document lint/type tooling in CLAUDE.md ... First repo-wide ruff format pass, plus a note in CLAUDE.md to run ruff and ty periodically. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>	2026-06-18 17:40:27 +02:00
lars	a867fc4aae	Add giant.analysis module for notebook-based model quality diagnostics ... Provides stratified marginal comparisons, joint-structure checks (correlation matrices, physically-coupled pairwise plots, direction alignment), and physical-constraint validation (unit-norm directions, non-negative raw targets) for a trained model's generated samples, building on the aggregate marginal/KL check already in giant.validate. Supports two entry points: live sampling against a checkpoint + val data (load_model_bundle/collect_samples), or loading a precomputed `giant predict --coord local` parquet directly (load_predicted_local) without needing the checkpoint at all. Predict output is now tagged with parquet schema metadata so the loader can verify a file's format and reject coord=global or untagged files with a clear error instead of guessing from column names. Also extends the config git-hash mismatch warning (added for --config loading) to checkpoint loading: both `giant predict` and analysis.load_model_bundle now look for a config.toml next to the checkpoint and warn (without failing) if it was generated from a different git commit. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>	2026-06-18 16:36:19 +02:00