Commit Graph

3 Commits

Author SHA1 Message Date
lars 9112e845e0 v0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite
CI / Lint (ruff check) (push) Successful in 27s
CI / Format (ruff format) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 23s
CI / Tests (push) Successful in 1m1s
CI / Lint (ruff check) (pull_request) Successful in 26s
CI / Format (ruff format) (pull_request) Successful in 27s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 27s
CI / Tests (pull_request) Successful in 1m2s
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 d656cf3109 Store a quantile grid instead of a raw reservoir sample in the setup cache
CI / Format (ruff format) (push) Successful in 26s
CI / Lint (ruff check) (push) Successful in 28s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 24s
CI / Tests (push) Successful in 58s
CI / Format (ruff format) (pull_request) Successful in 28s
CI / Lint (ruff check) (pull_request) Successful in 28s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 28s
CI / Tests (pull_request) Successful in 59s
NormalizerEntry.energy_reservoir_sample kept 100k raw energy values purely
to seed EnergyRouter centers via np.quantile at load time, which alone
accounted for most of the setup cache sidecar's ~2MB size (float32 values
round-tripped through Python floats serialize at full double precision).
Only a handful of quantile levels are ever read back, so collapse the
sample to a fixed 1001-point quantile grid at save time and interpolate
arbitrary levels from it at use time instead — about 100x smaller with
negligible (<0.001) error on the levels that matter. Bumps the cache
format version since old sidecars have no such grid to fall back on.
2026-07-30 13:32:46 +02:00
lars 471a81b5e7 Add dwarf warm-cache to precompute the setup-stage sidecar
CI / Format (ruff format) (push) Successful in 26s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 37s
CI / Lint (ruff check) (pull_request) Successful in 36s
CI / Format (ruff format) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 33s
CI / Tests (push) Successful in 1m43s
CI / Tests (pull_request) Successful in 1m39s
Lets the vocab maps, event-id split index, and normalizer stats be
warmed once for a dataset (right after `dwarf convert`, or before a
`dwarf hparam-scan` sweep) without needing to also start training.

Extracts the setup-stage logic out of giant/pipeline.py:run_train_job
into a standalone run_setup_stage() (returning a SetupStageResult),
reused by both run_train_job and the new dwarf command's
scripts/warm_setup_cache.py — a behavior-preserving refactor, covered
by the existing test_pipeline.py suite.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-30 10:53:31 +02:00