8 Commits

Author	SHA1	Message	Date
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	8019a80563	Refactor train.py into giant/training/ around a metrics collector ... Every metric name used to exist in four places: the dict keys each StageTrainer returned, the hardcoded _metrics_fields() column list, the ~110-line metrics_row assembly in train(), and the tqdm/summary formatting. The two had to be kept in exact correspondence by hand or csv.DictWriter would raise. Each metric is now declared once, as a MetricSpec on the trainer that computes it. MetricsCollector derives the CSV header and W&B payload from those declarations and owns all accumulation, so train() no longer carries a running sum, and every isinstance(tr, WGANStageTrainer) branch is gone — replaced by four trainer hooks (batch_loss, summary, val_objective, supports_val_loss). giant/train.py (1875 lines) becomes giant/training/: trainers.py StageSpec + shared StageTrainer base + the two subclasses metrics.py MetricSpec, MetricsCollector stage2_inputs.py the pure AR/teacher-forcing tensor helpers, moved verbatim loop.py train() (225 lines, was ~514) + graceful shutdown checkpoint.py build/load, lifted out of train()'s closures The trainers shared ~15 identical constructor arguments and copy-pasted their cosine-warmup lambda, EMA setup, state_dict/load_state_dict, resume_lr and train_mode/eval_mode. StageSpec resolves one stage's config once (constructors go from 24 and 22 keyword arguments to (spec, model, device)), the base class holds the rest, and build_stage_trainers drops from ~100 lines to 15. Metric columns are renamed to a uniform stage/split/metric scheme (stage1/train/loss, stage2/train/d_loss, stage1/lr, stage1/router/entropy, val/loss, ...). Old metrics.csv files and W&B history are not comparable. The checkpoint format is unchanged. BEHAVIOR CHANGE — WGAN best-checkpoint selection. The old code meant to score a WGAN stage on its marginal KL, but the guard `{n: kl for n in wgan_names if n not in val_loss_per_stage}` could never fire: val_loss_per_stage was pre-seeded with 0.0 for every stage, so a WGAN stage contributed a flat 0.0 and the KL was written to metrics.csv without ever influencing best.pt. val_objective now returns it as intended. On the test harness's default flow+wgan config val_loss went from 2.182 (stage 1 only) to 15.137 (stage 1 + KL 12.954), and which epoch won changed. Runs before this commit picked their best checkpoint on the non-adversarial stages alone. Written up in docs/v0.3.0-followups.md. Verified: 699 tests pass; ruff, ruff format and ty clean. Baseline-vs- refactor metrics.csv compared across five configs (flow+wgan, AR+onehot, routed, both-flow, AR-flow) — every comparable value bit-identical except val/loss where the fix applies. Resume appends without a duplicate header and reproduces a HEAD worktree's per-epoch losses and LRs exactly across the resume boundary. A refactored last.pt loads through cli.py:_load_model_weights in both raw and ema modes. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-07 17:03:20 +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	4fc15ecdfc	v0.3.0 step 4: type map + particle_type.target = "onehot"/"embedding" ... Builds the shared top-N-plus-other PDG/material maps (pooling both primary and secondary occurrences for PDG, directly targeting the meeting's species-collapse failure mode) and wires up conditioning.{particle,material} = "onehot" plus stage2_model.particle_type.target in ("onehot", "embedding") end-to-end: setup-cache persistence, Stage2OneShot's type_head (flow/ddpm) vs. folded+ST-Gumbel-relaxed adversarial slice (wgan), and the corresponding CE/MSE training losses. particle_type.target = "physical" stays byte-for-byte unchanged, keeping the v0.2 migration shim's bit-identical guarantee intact. giant predict/rollout fail loudly on a onehot/embedding checkpoint until full decode support lands in step 6. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-08-06 15:43:48 +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	8065df896e	Scope wandb run config to only-active hyperparameters ... Router-only knobs (lambda_balance/lambda_proc/lambda_entropy/ gumbel_tau_start/_end) and WGAN-only knobs (n_critic/gp_weight) were being logged to wandb's top-level run config unconditionally, even for runs where routing or WGAN mode is off, implying hyperparameters from an inactive code path. Extract _wandb_run_config and only include each group when its gate is actually true (router.enabled / mode=="wgan"); the full model_config (with its router sub-dict) is still always logged in full, so no information is lost. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-07-30 16:46:19 +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