lars
5eec4c250a
CI / Format (ruff format) (push) Successful in 25s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 29s
CI / Type check (ty) (push) Successful in 31s
CI / Format (ruff format) (pull_request) Successful in 37s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 38s
CI / Tests (push) Successful in 1m34s
CI / Tests (pull_request) Successful in 1m30s
Extends default_out_dir_name's non-default-field convention to the router's new gumbel combine-weight flag and its learnable-knob toggles, so gumbel sweep configs (learn_centers on/off, learn_width, learn_temperature) resolve to distinguishable checkpoint directory names instead of colliding on the same r-<type><n> token. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
315 lines
10 KiB
Python
315 lines
10 KiB
Python
from datetime import datetime
|
|
|
|
from giant import config as gconfig
|
|
|
|
|
|
def _write_config(path, git_hash):
|
|
path.write_text(
|
|
f"""
|
|
[train]
|
|
epochs = 5
|
|
|
|
[model]
|
|
hidden_dim = 64
|
|
|
|
[meta]
|
|
git_hash = "{git_hash}"
|
|
"""
|
|
)
|
|
|
|
|
|
def test_merge_cli_overrides_applies_file_then_cli(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
|
|
path = tmp_path / "config.toml"
|
|
_write_config(path, "abc123")
|
|
|
|
cfg = gconfig.merge_cli_overrides(
|
|
gconfig.DEFAULT_CONFIG,
|
|
path,
|
|
train_overrides={},
|
|
model_overrides={"hidden_dim": 128},
|
|
)
|
|
assert cfg["train"]["epochs"] == 5 # from file
|
|
assert cfg["model"]["hidden_dim"] == 128 # CLI override wins over file
|
|
|
|
|
|
def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, capsys):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
|
|
path = tmp_path / "config.toml"
|
|
_write_config(path, "old111")
|
|
|
|
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
|
|
|
assert cfg["train"]["epochs"] == 5 # does not fail, config still applied
|
|
captured = capsys.readouterr()
|
|
assert "warning" in captured.err
|
|
assert "old111" in captured.err
|
|
assert "current999" in captured.err
|
|
|
|
|
|
def test_merge_cli_overrides_no_warning_on_matching_git_hash(
|
|
tmp_path, monkeypatch, capsys
|
|
):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
|
|
path = tmp_path / "config.toml"
|
|
_write_config(path, "same123")
|
|
|
|
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
|
assert capsys.readouterr().err == ""
|
|
|
|
|
|
def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
|
|
tmp_path, monkeypatch, capsys
|
|
):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
|
|
path = tmp_path / "config.toml"
|
|
_write_config(path, "abc123")
|
|
|
|
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
|
assert capsys.readouterr().err == ""
|
|
|
|
|
|
def test_merge_cli_overrides_no_warning_when_meta_section_absent(
|
|
tmp_path, monkeypatch, capsys
|
|
):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
|
|
path = tmp_path / "config.toml"
|
|
path.write_text("[train]\nepochs = 5\n")
|
|
|
|
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
|
assert capsys.readouterr().err == ""
|
|
|
|
|
|
def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
|
|
tmp_path, monkeypatch, capsys
|
|
):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
|
|
ckpt_path = tmp_path / "best.pt"
|
|
ckpt_path.write_bytes(b"") # contents irrelevant, only its directory is used
|
|
_write_config(tmp_path / "config.toml", "old111")
|
|
|
|
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
|
|
|
|
captured = capsys.readouterr()
|
|
assert "warning" in captured.err
|
|
assert "old111" in captured.err
|
|
assert "current999" in captured.err
|
|
|
|
|
|
def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
|
|
tmp_path, monkeypatch, capsys
|
|
):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
|
|
ckpt_path = tmp_path / "best.pt"
|
|
ckpt_path.write_bytes(b"")
|
|
|
|
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
|
|
assert capsys.readouterr().err == ""
|
|
|
|
|
|
def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
|
|
tmp_path, monkeypatch, capsys
|
|
):
|
|
monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
|
|
ckpt_path = tmp_path / "best.pt"
|
|
ckpt_path.write_bytes(b"")
|
|
_write_config(tmp_path / "config.toml", "same123")
|
|
|
|
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
|
|
assert capsys.readouterr().err == ""
|
|
|
|
|
|
def test_resolve_expert_dims_default_config_inherits_hidden_dim_and_n_blocks():
|
|
# The unset sentinel (expert_hidden_dim/n_blocks == 0 in DEFAULT_CONFIG)
|
|
# is exactly the bug fixed by resolve_expert_dims: it must not silently
|
|
# fall back to some other hardcoded default, only to the monolith's own
|
|
# hidden_dim/n_blocks, so --hidden-dim/--n-blocks reach the experts too.
|
|
router_cfg = dict(gconfig.DEFAULT_CONFIG["model"]["router"])
|
|
assert router_cfg["expert_hidden_dim"] == 0
|
|
assert router_cfg["expert_n_blocks"] == 0
|
|
|
|
hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
|
|
assert (hidden_dim, n_blocks) == (512, 6)
|
|
|
|
|
|
def test_resolve_expert_dims_missing_keys_also_inherit():
|
|
hidden_dim, n_blocks = gconfig.resolve_expert_dims({}, 512, 6)
|
|
assert (hidden_dim, n_blocks) == (512, 6)
|
|
|
|
|
|
def test_default_config_gumbel_router_defaults_off():
|
|
# Straight-through Gumbel-softmax combine weights (giant.model.network.
|
|
# Router.combine_weights) must be opt-in — existing routed configs and
|
|
# checkpoints should be unaffected unless gumbel is explicitly enabled.
|
|
router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"]
|
|
assert router_cfg["gumbel"] is False
|
|
assert router_cfg["gumbel_tau_start"] == 1.0
|
|
assert router_cfg["gumbel_tau_end"] == 0.1
|
|
|
|
|
|
def test_resolve_expert_dims_explicit_override_wins():
|
|
router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3}
|
|
hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
|
|
assert (hidden_dim, n_blocks) == (128, 3)
|
|
|
|
|
|
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)
|
|
|
|
|
|
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_router_gumbel_shown_when_enabled():
|
|
cfg = _default_cfg(
|
|
router={"enabled": True, "type": "energy", "n_experts": 8, "gumbel": True}
|
|
)
|
|
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_gum"
|
|
|
|
|
|
def test_default_out_dir_name_router_gumbel_omitted_when_router_disabled():
|
|
cfg = _default_cfg(
|
|
router={"enabled": False, "type": "energy", "n_experts": 8, "gumbel": True}
|
|
)
|
|
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
|
|
|
|
|
|
def test_default_out_dir_name_router_learn_centers_shown_only_when_disabled():
|
|
cfg_default = _default_cfg(
|
|
router={"enabled": True, "type": "energy", "n_experts": 8}
|
|
)
|
|
assert (
|
|
gconfig.default_out_dir_name(cfg_default, now=_NOW) == "20260729_1430_r-energy8"
|
|
)
|
|
|
|
cfg_off = _default_cfg(
|
|
router={
|
|
"enabled": True,
|
|
"type": "energy",
|
|
"n_experts": 8,
|
|
"learn_centers": False,
|
|
}
|
|
)
|
|
assert (
|
|
gconfig.default_out_dir_name(cfg_off, now=_NOW)
|
|
== "20260729_1430_r-energy8_nolc"
|
|
)
|
|
|
|
|
|
def test_default_out_dir_name_router_learn_width_and_temperature_shown():
|
|
cfg = _default_cfg(
|
|
router={
|
|
"enabled": True,
|
|
"type": "energy",
|
|
"n_experts": 8,
|
|
"learn_width": True,
|
|
}
|
|
)
|
|
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_lw"
|
|
|
|
cfg2 = _default_cfg(
|
|
router={
|
|
"enabled": True,
|
|
"type": "energy",
|
|
"n_experts": 8,
|
|
"learn_temperature": True,
|
|
}
|
|
)
|
|
assert gconfig.default_out_dir_name(cfg2, now=_NOW) == "20260729_1430_r-energy8_lt"
|
|
|
|
|
|
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
|