7 Commits

Author	SHA1	Message	Date
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	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	a489991a3b	Document the differentiability position and its validation obligation ... The categorical type path is not differentiable, and v0.3.0 accepts that: the expected contribution of the broken path to the total gradient is assumed negligible. Record it as an assumption with an explicit obligation to demonstrate it, not as a settled result. Separates the three things "broken" covers, since they have different status: per-token loss under teacher forcing is fine (softmax CE needs no sampling); ST-Gumbel into the critic is biased rather than absent (hard forward, soft backward); full shower-rollout backprop was already structurally non-differentiable once secondaries branch, so the switch costs nothing that was not already lost. The accepted claim concerns only the middle one. Lists three ways to falsify it, cheapest first: gradient-magnitude accounting through the type slice vs the continuous slices, a detached-type ablation, and an estimator swap against REINFORCE if those are inconclusive. The first is wired into implementation step 5 so evidence accrues during the architecture comparison rather than in a dedicated run afterwards, and the fallback if the ratio is not small is a config change (target = "physical" or a non-adversarial CE head), not a redesign. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-06 09:57:44 +02:00
lars	376bdb9d08	Refine v0.3.0 design: defaults, deferred scope, open questions ... Config defaults: dropout 0.1 -> 0.0, wandb false -> true. The v0.3.0 work is a sequence of architecture comparisons, and an unlogged run is not comparable, so W&B is on unless explicitly disabled. Restructure the open-questions section into settled / deferred / tracked / still-open, since most of it is now decided: - other_policy three-way switch and separate flow/ddpm sub-tables are approved as specified. - stop_token is schema-valid but raises "not implemented in v0.3.0"; charge conservation gets no key at all and is left deliberately undesigned, to be worked out on its own terms rather than pre-shaped by this refactor. The speculative charge-mask sketch is removed. - Logging the L1 decode-distance distribution under target = "embedding" becomes tracked implementation work, landing with the rollout decode. - estimate_batch_size recalibration becomes implementation step 8, last: the activation-memory profile is not knowable until the AR trunk and history encoder are final, so it is measured on real hardware with the example configs rather than guessed. Only the differentiability question for Jan remains genuinely open. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-06 09:54:48 +02:00
lars	f390884f67	Add v0.3.0 design doc: Stage-2 autoregressive redesign ... Design contract for the v0.3.0 config break and network.py refactor, following the 2026-08-04 meeting with Jan. Not implemented yet. The 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photons, ~4M hallucinated -14 muon antineutrinos). The response pivots Stage 2 to autoregressive generation in descending-energy order with teacher forcing, and reverts the particle type to a categorical representation. That needs a config break: [conditioning] / [stage1_model] / [stage2_model] / [train] replace the single global train.mode and [model] block, so per-stage generators (stage1 flow + stage2 wgan), stage-2-only training, and one-shot-vs-autoregressive comparison all become expressible. The particle and material conditioning axes are configured independently and mix freely, each with physical / embedding / onehot modes; the stage-2 type target mirrors the same three names, with conditioning.particle.emb_dim sizing both so the two share one class map. network.py collapses from ten permutation classes (stage x objective x routed) into composable parts — encoder x trunk x objective — which also makes routed WGAN work for the first time; it was only ever rejected because no routed WGAN generator class existed. The doc specifies every config option, the v0.2 migration (shim for both configs and checkpoints, gated on a bit-identical output diff), the refactor, and the implementation order. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>	2026-08-04 15:25:32 +02:00
lars	b6893b0118	chore: remove stray CUDA sanity script and stale Phase 2 planning doc ... test-cuda.py was a one-off local CUDA check, not part of tests/ or scripts/. docs/phase2_plan.md is superseded by the "Phase 2 (implemented)" section of CLAUDE.md's roadmap. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>	2026-07-24 13:49:00 +02:00
lars	e6e0eb22bf	Implement Phase 2: secondary particle prediction ... Two-stage factorisation: Stage 1 predicts 9D primary kinematics + n_sec classification head (COND_DIM reduced to 8, dropping n_sec/e_sec inputs); Stage 2 (SecondaryDecoder) generates K_MAX=15 secondary slots via masked flow matching over (stick_logit, local_dir, type_emb) conditioned on Stage 1 output. Joint training with combined loss L_s1 + λ_nsec*L_nsec + λ_s2*L_s2. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>	2026-06-29 11:34:31 +02:00