29 Commits

Author	SHA1	Message	Date
lars	fce47b128c	Add class-balanced secondary particle-type loss (gitea #44) ... The v0.3.0 pivot exists because the 2026-08-03 WGAN rollout benchmark produced zero photon secondaries and ~4M hallucinated antineutrinos — even with a correctly-sized top-N species vocabulary (gitea #29), plain cross-entropy over a class distribution spanning orders of magnitude still under-predicts rare-but-physical species. stage2_model.particle_type.class_weighting = "none" | "inverse_freq" (default "none", fully back-compat) weights the stage-2 type head's CE loss (FlowDDPMStageTrainer._type_loss) by inverse class frequency, normalized to mean 1 so switching it on doesn't rescale the type loss against particle_type.lambda / the generator loss it's summed with. The per-class counts the weighting needs don't already exist despite the issue's premise: _topn_plus_other_map (giant/data/loader.py) previously kept counts only for keys folded into "other", dropping the kept classes' counts on the floor. TopNMap now carries class_counts (index -> count), round-tripped through the setup-cache sidecar (format version bumped 3->4, since existing sidecars have none) and through checkpoints (tolerantly — a pre-#44 checkpoint decodes to {}, since only training-time loss weighting reads it, not inference). Decisions made during planning (with the user): dropped "effective_num" from the issue's proposed three-way enum (no beta hyperparameter to design around) — final domain is "none" | "inverse_freq". Weights are mean-1-normalized. validate_config rejects class_weighting != "none" combined with particle_type.target != "onehot" or stage2_model.generator == "wgan" (both have no class CE to weight), following the #28/#30 dead-key-must-not-go-silent convention. Branch fix/issue-44 off master. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-17 23:02:50 +02:00
lars	78978769f6	Add bf16 autocast to the training loop (gitea #47) ... giant/ had no autocast/GradScaler/torch.compile anywhere despite the project's ~10x-native-Geant4 eval-budget target. This adds bf16 mixed precision to the training step (both FlowDDPMStageTrainer and WGANStageTrainer) via a new train.precision config key ("fp32" default, "bf16" opt-in) and giant.training.amp.resolve_autocast. torch.compile is a separate, much larger surface (data-dependent routed dispatch, the autoregressive sampler's per-token control flow, arbitrary rollout batch sizes) and is left for a follow-up issue, per discussion. Scope decisions made during planning: - fp32 + bf16 only, no fp16/GradScaler. fp16 breaks two things in this codebase: routers.py's three 1e-8 epsilons sit below fp16's ~6e-8 subnormal floor, and gradient_penalty's grad norm overflows fp16's range at ordinary early-WGAN-GP gradient magnitudes. Every training GPU in the fleet (A100/L40S/H200/RTX 4070) has native bf16; only pre-Ampere V100s would need fp16. - resolve_autocast raises loudly if bf16 is requested on hardware that can't do it, rather than silently falling back to fp32. - Autocast wraps the training step only; val_loss (and the best-checkpoint selection it drives) stays fp32 so it's comparable across every run recorded so far. - _route_forward's mixture accumulator (giant/model/trunks.py) was a hard-fp32 torch.zeros with no dtype, so under autocast a RoutedTrunk silently returned a different output dtype than an unrouted ExpertTrunk purely because router.enabled was set. Fixed to match the experts' own dtype; the gate weights (forced fp32 for their own numerical stability) are cast down before combining, so the mixture's numerics stay solid without reintroducing the dtype split. - Added explicit fp32 guards (autocast(enabled=False)) around spots that are correct in fp32 but degrade quietly rather than crash in bf16: the router's balance/entropy losses and gate softmax, the stage-2 stick-breaking cumprod, and gradient_penalty's double-backward + grad norm. Benchmarked on the local RTX 4070 against configs/baseline.toml's hyperparams (hidden_dim 512/6 blocks, bs 4096) on a synthetic dataset: bf16 gave 1.05-1.35x training throughput and 18-33% lower peak GPU memory across one-shot/routed/autoregressive stage-2 configs, with the autoregressive path (the dominant cost per baseline.toml) benefiting most on both axes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-17 15:36:12 +02:00
lars	87e37ebe14	Add per-stage init_from/freeze (gitea #42) ... stage{1,2}_model.active = false already trains one stage alone, but the checkpoint it writes holds only that stage, so giant rollout refuses it -- the "retrain stage 2 alone against a fixed, known-good stage 1" experiment the 2026-08-03 species failure calls for wasn't runnable end to end. Adds stage{1,2}_model.init_from (a checkpoint .pt to load this stage's weights from before training) and .freeze (never update them), symmetric across both stages. Both stages stay active = true, so both get built and both land in the output checkpoint -- the frozen stage is merely initialized from disk instead of from scratch. Decisions made during planning: - Soft freeze: forward/backward still run every batch (loss/grad_norm stay meaningful, no autograd special-casing), only optimizer.step() (and, for the frozen stage, lr_sched.step()/EMA update) is skipped -- weights are byte-identical for the whole run. This is StageTrainer._step_optimizer, shared by the non-adversarial path and both halves (generator + critic) of the WGAN path, so a frozen WGAN stage's critic freezes too. - validate_config requires init_from whenever freeze = true, unless the run is a --resume (a resumed frozen stage's weights come from the resume checkpoint instead) -- freezing a randomly-initialized model is almost certainly a mistake. - CLI flags on both `giant train` and `giant new-run` (--stage{1,2}-init-from/--stage{1,2}-freeze), matching every other per-stage model knob's existing treatment. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-17 14:23:20 +02:00
lars	48faaee79d	Implement stage2_model.stage1_context = "sampled" (gitea #41) ... Stage 2 was trained on ground-truth stage-1 outcomes but deployed on sampled ones, and in a rollout that gap compounds over every step of every track — the same train/inference gap teacher_forcing="scheduled" already closes within stage 2, just never applied at the stage boundary. "sampled" was declared in the schema but rejected loudly by validate_config as unimplemented; this lands the real implementation. Mirrors the existing scheduled-sampling precedent rather than a hard switch: new stage2_model.ctx_p_start/ctx_p_end (defaults 1.0 -> 0.0) linearly ramp P(condition on ground truth) from epoch 0 to the final epoch, so stage 2 doesn't chase a wildly moving stage-1 target early in training. Per the plan discussed with the user: the sample is drawn from stage 1's sampling_model() (EMA weights when present, matching what inference actually deploys), mixed per example via a Bernoulli draw (never blended within a row), and validation always uses the ground truth regardless of the schedule. Fixes a latent bug the same pattern would otherwise have hit: every sampler in giant/sample.py flips its model to .eval() with no restore, so sampling from the raw (non-EMA) stage-1 model mid-step now explicitly restores its .training flag afterward to avoid silently corrupting stage 1's own training mode for the rest of the epoch. validate_config now enforces stage1_context in {"truth", "sampled"}, requires both stages active for "sampled" (nothing to sample from otherwise), range-checks ctx_p_start/ctx_p_end, and rejects the ctx_p_start = ctx_p_end = 1.0 configuration as an unadvertised no-op identical to "truth". Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-17 12:27:13 +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	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	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