lars
9bf5874308
CI / Format (ruff format) (push) Successful in 31s
CI / Lint (ruff check) (push) Successful in 31s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 44s
CI / Type check (ty) (push) Successful in 46s
CI / Format (ruff format) (pull_request) Successful in 39s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 39s
CI / Tests (pull_request) Successful in 3m38s
CI / Tests (push) Successful in 3m45s
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>
76 lines
2.7 KiB
Python
76 lines
2.7 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 import config as gconfig
|
|
from giant.constants import K_MAX
|
|
from giant.pipeline import run_setup_stage
|
|
|
|
|
|
def run_warm_setup_cache(
|
|
data: str,
|
|
val_fraction: float = 0.1,
|
|
seed: int = 0,
|
|
particle_conditioning: str = "physical",
|
|
material_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`/`particle_conditioning`/`material_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. The two
|
|
conditioning axes are independent and may differ.
|
|
`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,
|
|
}
|
|
# Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so
|
|
# run_setup_stage always sees every key it might read (e.g.
|
|
# conditioning.particle.emb_dim, stage2_model.particle_type.target) at
|
|
# its real default, not silently missing/None — see issues.md Issue 1.
|
|
# 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 = gconfig.merge_cli_overrides(
|
|
gconfig.DEFAULT_CONFIG,
|
|
None,
|
|
{
|
|
"conditioning": {
|
|
"particle": {"type": particle_conditioning},
|
|
"material": {"type": material_conditioning},
|
|
},
|
|
"stage1_model": {"router": router_cfg},
|
|
"stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
|
|
},
|
|
)
|
|
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.")
|