lars
9112e845e0
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>
64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
"""dwarf warm-cache — precompute `giant train`'s setup-stage sidecar ahead of time.
|
|
|
|
Thin wrapper around `giant.pipeline.run_setup_stage` so a dataset's vocab
|
|
maps, event-id split index, and normalizer stats can be warmed once — e.g.
|
|
right after `dwarf convert`, or before kicking off a `dwarf hparam-scan`
|
|
sweep — without needing to also start training. See giant/data/setup_cache.py
|
|
for the sidecar itself.
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
from giant.pipeline import run_setup_stage
|
|
|
|
|
|
def run_warm_setup_cache(
|
|
data: str,
|
|
val_fraction: float = 0.1,
|
|
seed: int = 0,
|
|
conditioning: str = "physical",
|
|
router_enabled: bool = False,
|
|
router_type: str = "energy",
|
|
n_experts: int = 4,
|
|
rebuild: bool = False,
|
|
echo=print,
|
|
) -> None:
|
|
"""Populate (or refresh) the setup cache sidecar for `data`.
|
|
|
|
`val_fraction`/`seed`/`conditioning` select the normalizer cache entry
|
|
(`giant.data.setup_cache.normalizer_key`) — pass the same values a later
|
|
`giant train` invocation will use so it hits this warmed entry.
|
|
`router_enabled`/`router_type`/`n_experts` only matter for
|
|
`router_type == "process"` (warms that `n_experts`'s process map); the
|
|
energy-router quantile summary is always collected regardless, so a
|
|
later `--router-type energy` run never needs to rescan just to seed
|
|
centers.
|
|
"""
|
|
router_cfg = {
|
|
"enabled": router_enabled,
|
|
"type": router_type,
|
|
"n_experts": n_experts,
|
|
}
|
|
# A minimal v0.3 cfg — only the keys run_setup_stage actually reads
|
|
# (conditioning.particle.type, stage{1,2}_model.router). This CLI only
|
|
# ever configures one router (matching today's single --router-type
|
|
# flag), so it's placed on stage1_model; stage2_model's stays disabled.
|
|
cfg = {
|
|
"conditioning": {
|
|
"particle": {"type": conditioning},
|
|
"material": {"type": conditioning},
|
|
},
|
|
"stage1_model": {"router": router_cfg},
|
|
"stage2_model": {"router": {"enabled": False}},
|
|
}
|
|
run_setup_stage(
|
|
Path(data),
|
|
val_fraction=val_fraction,
|
|
seed=seed,
|
|
cfg=cfg,
|
|
cache_setup=True,
|
|
rebuild_setup_cache=rebuild,
|
|
echo=echo,
|
|
)
|
|
echo("setup cache warmed.")
|