lars
d25dfc0343
CI / Format (ruff format) (push) Successful in 40s
CI / Lint (ruff check) (push) Successful in 40s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 42s
CI / Type check (ty) (push) Successful in 44s
CI / Format (ruff format) (pull_request) Successful in 36s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 37s
CI / Tests (push) Successful in 3m48s
CI / Tests (pull_request) Successful in 3m44s
warm-cache built its config from DEFAULT_CONFIG with only a handful of flags overridable, so it had no way to express settings like stage2_model.particle_type.n_classes. configs/baseline.toml sets that to 32; warm-cache always warmed the pdg top-N map under the emb_dim default (16) instead, so a `giant train --config configs/baseline.toml` run silently missed the cache and repaid the full parquet scan warm-cache exists to avoid. warm-cache now accepts the same --config a training run takes and resolves every value run_setup_stage needs (val_fraction/seed, conditioning types, both stages' router, particle_type.n_classes, ...) from one gconfig.merge_cli_overrides + validate_config pass, exactly like giant train's own pipeline does — so warming and training are guaranteed to agree. Per user decision, --config is mutually exclusive with the individual --val-fraction/--seed/--particle-conditioning/ --material-conditioning/--router*/flags (rejected outright rather than silently layered on top), since a hardcoded CLI default clobbering an unset config value is the same failure mode one level down. Also drops a hardcoded stage2_model.router/k_max override that was a no-op against today's defaults but would have clobbered a config setting either one away from its default — same bug class. Adding validate_config surfaced that the existing test_warm_cache_router_process_warms_proc_map test was warming a router.type="process" + conditioning.particle.type="physical" (the CLI's old hardcoded default) combination that giant train's own validate_config would already reject as incompatible — fixed by passing --particle-conditioning embedding, which is what a working --router-type process run actually requires. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
99 lines
4.0 KiB
Python
99 lines
4.0 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.pipeline import run_setup_stage
|
|
|
|
|
|
def run_warm_setup_cache(
|
|
data: str,
|
|
config_path: Path | None = None,
|
|
val_fraction: float | None = None,
|
|
seed: int | None = None,
|
|
particle_conditioning: str | None = None,
|
|
material_conditioning: str | None = None,
|
|
router_enabled: bool | None = None,
|
|
router_type: str | None = None,
|
|
n_experts: int | None = None,
|
|
rebuild: bool = False,
|
|
echo=print,
|
|
) -> None:
|
|
"""Populate (or refresh) the setup cache sidecar for `data`.
|
|
|
|
Two mutually exclusive ways to select what to warm for (enforced by the
|
|
caller, `giant.tools.dwarf.warm_cache` — this function just trusts
|
|
whichever combination it's given):
|
|
|
|
- `config_path`: the same TOML `giant train --config` takes. Every value
|
|
`run_setup_stage` needs (`train.val_fraction`/`seed`,
|
|
`conditioning.particle`/`material.type`, both stages' `router`,
|
|
`stage2_model.particle_type.n_classes`, ...) is read from the one
|
|
resulting merged `cfg`, so a later `giant train --config <same file>`
|
|
run resolves to exactly the same cache keys — see gitea #59.
|
|
- The individual flags below: `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.
|
|
|
|
Any flag left `None` is omitted from the merge, so it falls back to
|
|
`DEFAULT_CONFIG`'s own value (or the config file's, if `config_path` is
|
|
given) instead of silently overriding it — see gitea #59.
|
|
"""
|
|
overrides: dict = {}
|
|
|
|
conditioning_overrides: dict = {}
|
|
if particle_conditioning is not None:
|
|
conditioning_overrides["particle"] = {"type": particle_conditioning}
|
|
if material_conditioning is not None:
|
|
conditioning_overrides["material"] = {"type": material_conditioning}
|
|
if conditioning_overrides:
|
|
overrides["conditioning"] = conditioning_overrides
|
|
|
|
# This CLI only ever configures one router (matching today's single
|
|
# --router-type flag), so it's placed on stage1_model; stage2_model's is
|
|
# left to DEFAULT_CONFIG/the config file rather than forced disabled.
|
|
router_overrides: dict = {}
|
|
if router_enabled is not None:
|
|
router_overrides["enabled"] = router_enabled
|
|
if router_type is not None:
|
|
router_overrides["type"] = router_type
|
|
if n_experts is not None:
|
|
router_overrides["n_experts"] = n_experts
|
|
if router_overrides:
|
|
overrides["stage1_model"] = {"router": router_overrides}
|
|
|
|
train_overrides: dict = {}
|
|
if val_fraction is not None:
|
|
train_overrides["val_fraction"] = val_fraction
|
|
if seed is not None:
|
|
train_overrides["seed"] = seed
|
|
if train_overrides:
|
|
overrides["train"] = train_overrides
|
|
|
|
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, overrides)
|
|
gconfig.validate_config(cfg)
|
|
run_setup_stage(
|
|
Path(data),
|
|
val_fraction=cfg["train"]["val_fraction"],
|
|
seed=cfg["train"]["seed"],
|
|
cfg=cfg,
|
|
cache_setup=True,
|
|
rebuild_setup_cache=rebuild,
|
|
echo=echo,
|
|
)
|
|
echo("setup cache warmed.")
|