Make default checkpoint out_dir name reflect only non-default hyperparams #19

Merged
lars merged 1 commits from feature/smart-outdir-naming into master 2026-07-29 13:38:16 +02:00
3 changed files with 198 additions and 11 deletions
+10 -10
View File
@@ -528,16 +528,16 @@ def train(
# hyperparam baked into the name (e.g. --lr for a fine-tune).
out_dir = resume.parent
else:
out_dir = Path(
f"checkpoints/{datetime.now().strftime('%Y%m%d_%H%M%S')}"
f"_{t['mode']}"
f"_h{m['hidden_dim']}"
f"_b{m['n_blocks']}"
f"_e{m['emb_dim']}"
f"_c{m['conditioning']}"
f"_lr{t['lr']}"
f"_bs{t['batch_size']}"
)
# Name only encodes what's non-default (see default_out_dir_name), so
# two runs with identical hyperparams in the same to-the-minute
# timestamp would otherwise collide on this name — which also
# doubles as the W&B run id (giant.train) — hence the suffix loop.
base_name = gconfig.default_out_dir_name(cfg)
out_dir = Path("checkpoints") / base_name
suffix = 2
while out_dir.exists():
out_dir = Path("checkpoints") / f"{base_name}_{suffix}"
suffix += 1
typer.echo(f"device: {_device}")
typer.echo(f"out_dir: {out_dir}")
+88
View File
@@ -1,3 +1,4 @@
import hashlib
import random
import subprocess
import sys
@@ -284,6 +285,93 @@ def resolve_expert_dims(
return expert_hidden_dim, expert_n_blocks
_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb"}
# Priority-ordered candidate fields for default_out_dir_name: (label, getter,
# formatter). `getter(train, model)` returns None when the field is at its
# default (and so should be omitted); otherwise formatter(value) renders the
# name token. The router is a single unit gated on `router.enabled` rather
# than one candidate per router key, since its type/n_experts are meaningless
# while disabled.
def _mode_candidate(train, model):
return None if train["mode"] == DEFAULT_CONFIG["train"]["mode"] else train["mode"]
def _router_candidate(train, model):
router = model["router"]
if router["enabled"] == DEFAULT_CONFIG["model"]["router"]["enabled"]:
return None
return f"r-{router['type']}{router['n_experts']}"
def _conditioning_candidate(train, model):
if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]:
return None
code = _CONDITIONING_CODE.get(model["conditioning"], model["conditioning"])
return f"c{code}"
def _default_field_candidate(section_key, field, prefix):
def _candidate(train, model):
section = train if section_key == "train" else model
value = section[field]
if value == DEFAULT_CONFIG[section_key][field]:
return None
return f"{prefix}{value}"
return _candidate
_OUT_DIR_NAME_CANDIDATES = [
("mode", _mode_candidate),
("router", _router_candidate),
("conditioning", _conditioning_candidate),
("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")),
("n_blocks", _default_field_candidate("model", "n_blocks", "b")),
("emb_dim", _default_field_candidate("model", "emb_dim", "e")),
("lr", _default_field_candidate("train", "lr", "lr")),
("batch_size", _default_field_candidate("train", "batch_size", "bs")),
("seed", _default_field_candidate("train", "seed", "seed")),
("epochs", _default_field_candidate("train", "epochs", "ep")),
]
_OUT_DIR_NAME_MAX_FIELDS = 6
def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
"""Build a default checkpoint out_dir name from what's non-default in `cfg`.
Only fields that differ from DEFAULT_CONFIG are included, so a fully
default run's name is just its timestamp — see
`_OUT_DIR_NAME_CANDIDATES` for the fixed, priority-ordered field list.
Beyond `_OUT_DIR_NAME_MAX_FIELDS` non-default fields, the remainder
collapse into a short deterministic hash suffix rather than growing the
name unboundedly. This name doubles as the run's W&B id (see
giant.train), which is the reason a timestamp is always included.
"""
now = now or datetime.now()
train, model = cfg["train"], cfg["model"]
tokens = []
overflow = []
for label, candidate in _OUT_DIR_NAME_CANDIDATES:
token = candidate(train, model)
if token is None:
continue
if len(tokens) < _OUT_DIR_NAME_MAX_FIELDS:
tokens.append(token)
else:
overflow.append(f"{label}={token}")
name = now.strftime("%Y%m%d_%H%M")
if tokens:
name += "_" + "_".join(tokens)
if overflow:
digest = hashlib.md5("|".join(sorted(overflow)).encode()).hexdigest()[:6]
name += f"_+{len(overflow)}more-{digest}"
return name
def seed_everything(seed: int) -> None:
random.seed(seed)
np.random.seed(seed)
+100 -1
View File
@@ -1,3 +1,5 @@
from datetime import datetime
from giant import config as gconfig
@@ -144,4 +146,101 @@ def test_resolve_expert_dims_explicit_override_wins():
def test_resolve_expert_dims_partial_override_mixes_explicit_and_inherited():
router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 0}
hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
assert (hidden_dim, n_blocks) == (128, 6) # n_blocks inherited, hidden_dim not
assert (hidden_dim, n_blocks) == (128, 6)
def _default_cfg(**overrides):
train_overrides = {
k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["train"]
}
model_overrides = {
k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["model"]
}
router_overrides = overrides.get("router")
if router_overrides:
model_overrides["router"] = router_overrides
return gconfig.merge_cli_overrides(
gconfig.DEFAULT_CONFIG, None, train_overrides, model_overrides
)
_NOW = datetime(2026, 7, 29, 14, 30)
def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
cfg = _default_cfg()
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
def test_default_out_dir_name_single_non_default_field():
cfg = _default_cfg(hidden_dim=512)
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_h512"
def test_default_out_dir_name_conditioning_embedding_shown_abbreviated():
cfg = _default_cfg(conditioning="embedding")
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb"
def test_default_out_dir_name_conditioning_default_omitted():
cfg = _default_cfg(conditioning="physical")
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
def test_default_out_dir_name_router_enabled_shown_as_unit():
cfg = _default_cfg(router={"enabled": True, "type": "energy", "n_experts": 8})
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8"
def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault():
cfg = _default_cfg(router={"enabled": False, "type": "pdg", "n_experts": 8})
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
def test_default_out_dir_name_mode_shown_bare_no_prefix():
cfg = _default_cfg(mode="wgan")
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan"
def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
cfg = _default_cfg(
mode="wgan",
router={"enabled": True, "type": "energy", "n_experts": 8},
conditioning="embedding",
hidden_dim=512,
n_blocks=8,
emb_dim=32,
lr=1e-3,
batch_size=2048,
seed=3,
epochs=200,
)
name = gconfig.default_out_dir_name(cfg, now=_NOW)
# First 6 by priority: mode, router, conditioning, hidden_dim, n_blocks, emb_dim.
assert name.startswith("20260729_1430_wgan_r-energy8_cemb_h512_b8_e32_+4more-")
digest = name.split("-")[-1]
assert len(digest) == 6
def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
base = dict(
mode="wgan",
router={"enabled": True, "type": "energy", "n_experts": 8},
conditioning="embedding",
hidden_dim=512,
n_blocks=8,
emb_dim=32,
lr=1e-3,
batch_size=2048,
seed=3,
epochs=200,
)
name_a = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW)
name_b = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW)
assert name_a == name_b # stable across calls with the same overflow set
changed = dict(base, epochs=999)
name_c = gconfig.default_out_dir_name(_default_cfg(**changed), now=_NOW)
assert (
name_c != name_a
) # differs when an overflowed value changes # n_blocks inherited, hidden_dim not