v0.3.0 step 1: new nested config schema, v0.2 migration shim
Replace the single global train.mode + [model] block with the four top-level blocks docs/v0.3.0-design.md specifies ([conditioning], [stage1_model], [stage2_model], [train]), so Stage 1 and Stage 2 can run independent generative objectives and Stage 2 can train standalone. - migrate_config translates old config.toml/checkpoint dicts on load, so nothing on /ceph goes dead; loudly rejects non-zero expert_hidden_dim/expert_n_blocks, which v0.3.0 no longer supports. - merge_cli_overrides/save_config generalize from one hardcoded nesting level (model.router) to arbitrary recursive depth. - default_out_dir_name candidates move to dotted paths against the new schema, with per-stage router/generator discriminators. - validate_config adds cross-block checks the per-block schema can't express (particle_type.target=embedding needs a matching conditioning mode, tie_to_stage1 needs an active stage 1, etc). - resolve_expert_dims is deleted (experts always inherit the stage's hidden_dim/n_res_blocks now) — pipeline.py/cli.py callers are left dangling on purpose, to be updated in the network.py/train.py steps that follow. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Claude Sonnet 5
+575
-222
@@ -1,46 +1,314 @@
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
from giant import config as gconfig
|
||||
|
||||
_CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs"
|
||||
|
||||
def _write_config(path, git_hash):
|
||||
path.write_text(
|
||||
f"""
|
||||
[train]
|
||||
epochs = 5
|
||||
|
||||
[model]
|
||||
hidden_dim = 64
|
||||
def _write_toml(path, git_hash=None, extra=""):
|
||||
meta = f'\n[meta]\ngit_hash = "{git_hash}"\n' if git_hash is not None else ""
|
||||
path.write_text(extra + meta)
|
||||
|
||||
[meta]
|
||||
git_hash = "{git_hash}"
|
||||
"""
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Conditioning enum
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_conditioning_enum_has_onehot():
|
||||
assert gconfig.Conditioning.onehot == "onehot"
|
||||
assert {c.value for c in gconfig.Conditioning} == {
|
||||
"physical",
|
||||
"embedding",
|
||||
"onehot",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _deep_merge
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_deep_merge_leaf_override_keeps_untouched_siblings():
|
||||
base = {"a": 1, "b": {"c": 2, "d": 3}}
|
||||
result = gconfig._deep_merge(base, {"b": {"c": 99}})
|
||||
assert result == {"a": 1, "b": {"c": 99, "d": 3}}
|
||||
|
||||
|
||||
def test_deep_merge_recurses_at_multiple_levels():
|
||||
base = {
|
||||
"stage1_model": {
|
||||
"hidden_dim": 256,
|
||||
"router": {"enabled": False, "type": "energy", "n_experts": 4},
|
||||
}
|
||||
}
|
||||
result = gconfig._deep_merge(base, {"stage1_model": {"router": {"enabled": True}}})
|
||||
assert result["stage1_model"]["hidden_dim"] == 256
|
||||
assert result["stage1_model"]["router"] == {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 4,
|
||||
}
|
||||
|
||||
|
||||
def test_deep_merge_does_not_mutate_base():
|
||||
base = {"a": {"b": 1}}
|
||||
gconfig._deep_merge(base, {"a": {"b": 2}})
|
||||
assert base == {"a": {"b": 1}}
|
||||
|
||||
|
||||
def test_deep_merge_non_dict_override_replaces_wholesale():
|
||||
base = {"a": {"b": 1}}
|
||||
result = gconfig._deep_merge(base, {"a": 5})
|
||||
assert result == {"a": 5}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# migrate_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_migrate_config_already_v3_returned_unchanged():
|
||||
cfg = {"meta": {"config_version": 3}, "stage1_model": {"generator": "flow"}}
|
||||
result = gconfig.migrate_config(cfg)
|
||||
assert result == cfg
|
||||
result["stage1_model"]["generator"] = "wgan"
|
||||
assert cfg["stage1_model"]["generator"] == "flow" # deep-copied, not aliased
|
||||
|
||||
|
||||
def test_migrate_config_empty_dict_still_injects_hardcoded_v02_facts():
|
||||
# No [train]/[model] at all still counts as "v0.2" (config_version
|
||||
# absent) — the hardcoded architectural facts fire unconditionally.
|
||||
new = gconfig.migrate_config({})
|
||||
assert "train" not in new
|
||||
assert new["conditioning"]["out_dim"] == 128
|
||||
assert new["conditioning"]["particle"]["n_layers"] == 2
|
||||
assert new["conditioning"]["material"]["n_layers"] == 2
|
||||
assert new["stage1_model"]["active"] is True
|
||||
assert new["stage1_model"]["flow"]["time_dim"] == 64
|
||||
assert new["stage1_model"]["ddpm"]["time_dim"] == 64
|
||||
assert new["stage2_model"]["active"] is True
|
||||
assert new["stage2_model"]["flow"]["time_dim"] == 64
|
||||
assert new["stage2_model"]["ddpm"]["time_dim"] == 64
|
||||
assert new["stage2_model"]["context_dim"] == 64
|
||||
assert new["stage2_model"]["decoder"] == "one_shot"
|
||||
assert new["stage2_model"]["particle_type"]["target"] == "physical"
|
||||
assert new["meta"] == {"config_version": 3}
|
||||
|
||||
|
||||
def test_migrate_config_mode_maps_to_both_stage_generators():
|
||||
new = gconfig.migrate_config({"train": {"mode": "wgan"}})
|
||||
assert new["stage1_model"]["generator"] == "wgan"
|
||||
assert new["stage2_model"]["generator"] == "wgan"
|
||||
|
||||
|
||||
def test_migrate_config_lambda_nsec_and_lambda_s2():
|
||||
new = gconfig.migrate_config({"train": {"lambda_nsec": 0.2, "lambda_s2": 2.0}})
|
||||
assert new["stage2_model"]["n_sec"]["lambda"] == 0.2
|
||||
assert new["stage2_model"]["lambda"] == 2.0
|
||||
|
||||
|
||||
def test_migrate_config_wgan_knobs_map_to_both_stages():
|
||||
new = gconfig.migrate_config(
|
||||
{"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}}
|
||||
)
|
||||
for stage in ("stage1_model", "stage2_model"):
|
||||
assert new[stage]["wgan"]["n_critic"] == 3
|
||||
assert new[stage]["wgan"]["gp_weight"] == 5.0
|
||||
assert new[stage]["wgan"]["critic_lr"] == 1e-4
|
||||
|
||||
|
||||
def test_merge_cli_overrides_applies_file_then_cli(tmp_path, monkeypatch):
|
||||
def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
|
||||
new = gconfig.migrate_config(
|
||||
{"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}}
|
||||
)
|
||||
for stage in ("stage1_model", "stage2_model"):
|
||||
assert new[stage]["hidden_dim"] == 128
|
||||
assert new[stage]["n_res_blocks"] == 4
|
||||
assert new[stage]["dropout"] == 0.2
|
||||
|
||||
|
||||
def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes():
|
||||
new = gconfig.migrate_config(
|
||||
{"model": {"emb_dim": 32, "conditioning": "embedding"}}
|
||||
)
|
||||
for axis in ("particle", "material"):
|
||||
assert new["conditioning"][axis]["emb_dim"] == 32
|
||||
assert new["conditioning"][axis]["type"] == "embedding"
|
||||
|
||||
|
||||
def test_migrate_config_noise_dim_maps_to_both_stages_wgan():
|
||||
new = gconfig.migrate_config({"model": {"noise_dim": 128}})
|
||||
assert new["stage1_model"]["wgan"]["noise_dim"] == 128
|
||||
assert new["stage2_model"]["wgan"]["noise_dim"] == 128
|
||||
|
||||
|
||||
def test_migrate_config_k_max_maps_to_stage2_only():
|
||||
new = gconfig.migrate_config({"model": {"k_max": 20}})
|
||||
assert new["stage2_model"]["k_max"] == 20
|
||||
assert "k_max" not in new.get("stage1_model", {})
|
||||
|
||||
|
||||
def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false():
|
||||
new = gconfig.migrate_config(
|
||||
{
|
||||
"model": {
|
||||
"router": {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 10,
|
||||
"temperature": 0.05,
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
assert new["stage1_model"]["router"] == {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 10,
|
||||
"temperature": 0.05,
|
||||
}
|
||||
assert new["stage2_model"]["router"] == {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 10,
|
||||
"temperature": 0.05,
|
||||
"tie_to_stage1": False,
|
||||
}
|
||||
|
||||
|
||||
def test_migrate_config_router_nonzero_expert_dims_raises():
|
||||
cfg = {
|
||||
"model": {
|
||||
"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}
|
||||
}
|
||||
}
|
||||
try:
|
||||
gconfig.migrate_config(cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "expert_hidden_dim" in str(e)
|
||||
|
||||
|
||||
def test_migrate_config_router_zero_expert_dims_dropped_silently():
|
||||
new = gconfig.migrate_config(
|
||||
{
|
||||
"model": {
|
||||
"router": {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 4,
|
||||
"expert_hidden_dim": 0,
|
||||
"expert_n_blocks": 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
)
|
||||
assert "expert_hidden_dim" not in new["stage1_model"]["router"]
|
||||
assert "expert_n_blocks" not in new["stage1_model"]["router"]
|
||||
|
||||
|
||||
def test_migrate_config_train_passthrough_is_exact():
|
||||
new = gconfig.migrate_config({"train": {"epochs": 7, "batch_size": 999, "seed": 3}})
|
||||
assert new["train"] == {"epochs": 7, "batch_size": 999, "seed": 3}
|
||||
|
||||
|
||||
def test_migrate_config_preserves_meta_git_hash():
|
||||
new = gconfig.migrate_config({"meta": {"git_hash": "abc123"}})
|
||||
assert new["meta"] == {"git_hash": "abc123", "config_version": 3}
|
||||
|
||||
|
||||
def test_migrate_config_real_router_fixture_raises_on_nonzero_expert_dims():
|
||||
cfg = gconfig.load_toml(_CONFIGS_DIR / "router_energy_n10_embedding.toml")
|
||||
try:
|
||||
gconfig.migrate_config(cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "expert_hidden_dim" in str(e)
|
||||
|
||||
|
||||
def test_migrate_config_real_wgan_fixture():
|
||||
cfg = gconfig.load_toml(_CONFIGS_DIR / "wgan_h128_b4_physical.toml")
|
||||
new = gconfig.migrate_config(cfg)
|
||||
assert new["stage1_model"]["generator"] == "wgan"
|
||||
assert new["stage2_model"]["generator"] == "wgan"
|
||||
for stage in ("stage1_model", "stage2_model"):
|
||||
assert new[stage]["hidden_dim"] == 128
|
||||
assert new[stage]["n_res_blocks"] == 4
|
||||
assert new[stage]["dropout"] == 0.0
|
||||
assert new["conditioning"]["particle"]["type"] == "physical"
|
||||
assert new["conditioning"]["material"]["type"] == "physical"
|
||||
assert new["train"]["epochs"] == 30
|
||||
assert new["train"]["warmup_epochs"] == 3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# merge_cli_overrides
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_merge_cli_overrides_defaults_only_matches_default_config():
|
||||
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {})
|
||||
assert cfg == gconfig.DEFAULT_CONFIG
|
||||
assert cfg is not gconfig.DEFAULT_CONFIG
|
||||
|
||||
|
||||
def test_merge_cli_overrides_nested_override_keeps_siblings():
|
||||
cfg = gconfig.merge_cli_overrides(
|
||||
gconfig.DEFAULT_CONFIG,
|
||||
None,
|
||||
{"stage1_model": {"router": {"enabled": True}}},
|
||||
)
|
||||
assert cfg["stage1_model"]["router"]["enabled"] is True
|
||||
assert cfg["stage1_model"]["router"]["type"] == "energy" # default preserved
|
||||
assert cfg["stage1_model"]["hidden_dim"] == 256 # untouched sibling section
|
||||
|
||||
|
||||
def test_merge_cli_overrides_file_then_explicit_override_precedence(
|
||||
tmp_path, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
|
||||
path = tmp_path / "config.toml"
|
||||
_write_config(path, "abc123")
|
||||
_write_toml(
|
||||
path,
|
||||
git_hash="abc123",
|
||||
extra="[train]\nepochs = 5\n\n[model]\nhidden_dim = 64\n",
|
||||
)
|
||||
|
||||
cfg = gconfig.merge_cli_overrides(
|
||||
gconfig.DEFAULT_CONFIG,
|
||||
path,
|
||||
train_overrides={},
|
||||
model_overrides={"hidden_dim": 128},
|
||||
{"stage1_model": {"hidden_dim": 128}},
|
||||
)
|
||||
assert cfg["train"]["epochs"] == 5 # from file
|
||||
assert cfg["model"]["hidden_dim"] == 128 # CLI override wins over file
|
||||
assert cfg["stage1_model"]["hidden_dim"] == 128 # explicit override wins over file
|
||||
assert cfg["stage2_model"]["hidden_dim"] == 64 # migrated from file, not overridden
|
||||
|
||||
|
||||
def test_merge_cli_overrides_migrates_v2_file_transparently(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
|
||||
path = tmp_path / "config.toml"
|
||||
_write_toml(
|
||||
path,
|
||||
git_hash="abc123",
|
||||
extra='[train]\nmode = "wgan"\n\n[model]\nconditioning = "embedding"\n',
|
||||
)
|
||||
|
||||
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
|
||||
assert cfg["stage1_model"]["generator"] == "wgan"
|
||||
assert cfg["stage2_model"]["generator"] == "wgan"
|
||||
assert cfg["conditioning"]["particle"]["type"] == "embedding"
|
||||
# hardcoded v0.2 fact still applied even though it's not a CLI-settable key
|
||||
assert cfg["conditioning"]["particle"]["n_layers"] == 2
|
||||
|
||||
|
||||
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")
|
||||
_write_toml(path, git_hash="old111", extra="[train]\nepochs = 5\n")
|
||||
|
||||
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
||||
|
||||
assert cfg["train"]["epochs"] == 5 # does not fail, config still applied
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
|
||||
captured = capsys.readouterr()
|
||||
assert "warning" in captured.err
|
||||
assert "old111" in captured.err
|
||||
@@ -52,9 +320,9 @@ def test_merge_cli_overrides_no_warning_on_matching_git_hash(
|
||||
):
|
||||
monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
|
||||
path = tmp_path / "config.toml"
|
||||
_write_config(path, "same123")
|
||||
_write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n")
|
||||
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
|
||||
assert capsys.readouterr().err == ""
|
||||
|
||||
|
||||
@@ -63,9 +331,9 @@ def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
|
||||
):
|
||||
monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
|
||||
path = tmp_path / "config.toml"
|
||||
_write_config(path, "abc123")
|
||||
_write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n")
|
||||
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
|
||||
assert capsys.readouterr().err == ""
|
||||
|
||||
|
||||
@@ -76,17 +344,295 @@ def test_merge_cli_overrides_no_warning_when_meta_section_absent(
|
||||
path = tmp_path / "config.toml"
|
||||
path.write_text("[train]\nepochs = 5\n")
|
||||
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {})
|
||||
gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
|
||||
assert capsys.readouterr().err == ""
|
||||
|
||||
|
||||
def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243"
|
||||
)
|
||||
cfg = gconfig.merge_cli_overrides(
|
||||
gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {}
|
||||
)
|
||||
assert cfg["stage1_model"]["generator"] == "flow"
|
||||
assert cfg["stage1_model"]["hidden_dim"] == 256
|
||||
assert cfg["stage2_model"]["hidden_dim"] == 256
|
||||
assert cfg["conditioning"]["particle"]["emb_dim"] == 16
|
||||
assert cfg["conditioning"]["particle"]["n_layers"] == 2 # migrated hardcoded fact
|
||||
assert cfg["train"]["epochs"] == 100
|
||||
assert cfg["stage2_model"]["decoder"] == "one_shot"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# save_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_save_config_round_trips_multi_level_nesting(tmp_path):
|
||||
cfg = {
|
||||
"stage1_model": {
|
||||
"hidden_dim": 256,
|
||||
"router": {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 4,
|
||||
},
|
||||
},
|
||||
"train": {"epochs": 100},
|
||||
}
|
||||
meta = {"config_version": 3, "git_hash": "abc123"}
|
||||
|
||||
gconfig.save_config(cfg, tmp_path, meta)
|
||||
loaded = gconfig.load_toml(tmp_path / "config.toml")
|
||||
|
||||
assert loaded["stage1_model"]["hidden_dim"] == 256
|
||||
assert loaded["stage1_model"]["router"] == {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 4,
|
||||
}
|
||||
assert loaded["train"] == {"epochs": 100}
|
||||
assert loaded["meta"] == meta
|
||||
|
||||
|
||||
def test_save_config_round_trips_three_level_nesting(tmp_path):
|
||||
cfg = {
|
||||
"stage2_model": {
|
||||
"decoder": "autoregressive",
|
||||
"n_sec": {"mode": "head", "lambda": 0.1},
|
||||
"router": {"tie_to_stage1": True},
|
||||
}
|
||||
}
|
||||
gconfig.save_config(cfg, tmp_path, {"config_version": 3})
|
||||
loaded = gconfig.load_toml(tmp_path / "config.toml")
|
||||
|
||||
assert loaded["stage2_model"]["decoder"] == "autoregressive"
|
||||
assert loaded["stage2_model"]["n_sec"] == {"mode": "head", "lambda": 0.1}
|
||||
assert loaded["stage2_model"]["router"] == {"tie_to_stage1": True}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# default_out_dir_name
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_NOW = datetime(2026, 7, 29, 14, 30)
|
||||
|
||||
|
||||
def _cfg_with(**dotted_overrides):
|
||||
"""Build a full DEFAULT_CONFIG-shaped dict with dotted-path overrides
|
||||
applied via _deep_merge, e.g. _cfg_with(**{"stage1_model.hidden_dim": 512})."""
|
||||
overrides: dict = {}
|
||||
for dotted, value in dotted_overrides.items():
|
||||
gconfig._set_path(overrides, dotted, value)
|
||||
return gconfig._deep_merge(gconfig.DEFAULT_CONFIG, overrides)
|
||||
|
||||
|
||||
def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
|
||||
assert (
|
||||
gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW)
|
||||
== "20260729_1430"
|
||||
)
|
||||
|
||||
|
||||
def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix():
|
||||
cfg = _cfg_with(**{"stage1_model.generator": "wgan"})
|
||||
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan"
|
||||
|
||||
|
||||
def test_default_out_dir_name_stage2_decoder_shown():
|
||||
cfg = _cfg_with(**{"stage2_model.decoder": "one_shot"})
|
||||
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_dec-one_shot"
|
||||
|
||||
|
||||
def test_default_out_dir_name_particle_type_target_shown():
|
||||
cfg = _cfg_with(**{"stage2_model.particle_type.target": "physical"})
|
||||
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_pt-physical"
|
||||
|
||||
|
||||
def test_default_out_dir_name_particle_conditioning_embedding_abbreviated():
|
||||
cfg = _cfg_with(**{"conditioning.particle.type": "embedding"})
|
||||
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb"
|
||||
|
||||
|
||||
def test_default_out_dir_name_stage1_router_shown_as_unit():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage1_model.router.enabled": True,
|
||||
"stage1_model.router.type": "energy",
|
||||
"stage1_model.router.n_experts": 8,
|
||||
}
|
||||
)
|
||||
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8"
|
||||
|
||||
|
||||
def test_default_out_dir_name_stage2_router_shown_as_unit_distinct_from_stage1():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage2_model.router.enabled": True,
|
||||
"stage2_model.router.type": "pdg",
|
||||
"stage2_model.router.n_experts": 3,
|
||||
}
|
||||
)
|
||||
assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s2r-pdg3"
|
||||
|
||||
|
||||
def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage1_model.router.enabled": False,
|
||||
"stage1_model.router.type": "pdg",
|
||||
"stage1_model.router.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 = _cfg_with(
|
||||
**{
|
||||
"stage1_model.router.enabled": True,
|
||||
"stage1_model.router.type": "energy",
|
||||
"stage1_model.router.n_experts": 8,
|
||||
"stage1_model.router.gumbel": True,
|
||||
}
|
||||
)
|
||||
assert (
|
||||
gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
|
||||
)
|
||||
|
||||
|
||||
def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage1_model.generator": "wgan",
|
||||
"stage2_model.generator": "flow",
|
||||
"stage2_model.decoder": "one_shot",
|
||||
"stage2_model.autoregressive.history": "attention",
|
||||
"stage2_model.particle_type.target": "physical",
|
||||
"stage1_model.router.enabled": True,
|
||||
"stage1_model.router.type": "energy",
|
||||
"stage1_model.router.n_experts": 8,
|
||||
"stage2_model.router.enabled": True,
|
||||
"stage2_model.router.type": "pdg",
|
||||
"stage2_model.router.n_experts": 3,
|
||||
}
|
||||
)
|
||||
name = gconfig.default_out_dir_name(cfg, now=_NOW)
|
||||
# First 6 by priority: stage1_generator, stage2_generator, stage2_decoder,
|
||||
# stage2_history, particle_type_target, stage1_router — stage2_router
|
||||
# overflows into the hash suffix.
|
||||
assert name.startswith(
|
||||
"20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+"
|
||||
)
|
||||
|
||||
|
||||
def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
|
||||
overrides = {
|
||||
"stage1_model.generator": "wgan",
|
||||
"stage2_model.generator": "flow",
|
||||
"stage2_model.decoder": "one_shot",
|
||||
"stage2_model.autoregressive.history": "attention",
|
||||
"stage2_model.particle_type.target": "physical",
|
||||
"stage1_model.router.enabled": True,
|
||||
"stage1_model.router.type": "energy",
|
||||
"stage1_model.router.n_experts": 8,
|
||||
"train.seed": 3,
|
||||
}
|
||||
name_a = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW)
|
||||
name_b = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW)
|
||||
assert name_a == name_b
|
||||
|
||||
changed = dict(overrides, **{"train.seed": 99})
|
||||
name_c = gconfig.default_out_dir_name(_cfg_with(**changed), now=_NOW)
|
||||
assert name_c != name_a
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# validate_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_validate_config_default_config_passes():
|
||||
gconfig.validate_config(gconfig.DEFAULT_CONFIG) # must not raise
|
||||
|
||||
|
||||
def test_validate_config_embedding_target_requires_embedding_conditioning():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage2_model.particle_type.target": "embedding",
|
||||
"conditioning.particle.type": "physical",
|
||||
}
|
||||
)
|
||||
try:
|
||||
gconfig.validate_config(cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "embedding" in str(e)
|
||||
|
||||
|
||||
def test_validate_config_embedding_target_passes_with_embedding_conditioning():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage2_model.particle_type.target": "embedding",
|
||||
"conditioning.particle.type": "embedding",
|
||||
}
|
||||
)
|
||||
gconfig.validate_config(cfg) # must not raise
|
||||
|
||||
|
||||
def test_validate_config_pdg_router_incompatible_with_physical_conditioning():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage1_model.router.enabled": True,
|
||||
"stage1_model.router.type": "pdg",
|
||||
"conditioning.particle.type": "physical",
|
||||
}
|
||||
)
|
||||
try:
|
||||
gconfig.validate_config(cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "pdg" in str(e)
|
||||
|
||||
|
||||
def test_validate_config_tie_to_stage1_requires_stage1_active():
|
||||
cfg = _cfg_with(
|
||||
**{
|
||||
"stage2_model.router.tie_to_stage1": True,
|
||||
"stage1_model.active": False,
|
||||
}
|
||||
)
|
||||
try:
|
||||
gconfig.validate_config(cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "tie_to_stage1" in str(e)
|
||||
|
||||
|
||||
def test_validate_config_stop_token_not_implemented():
|
||||
cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"})
|
||||
try:
|
||||
gconfig.validate_config(cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "stop_token" in str(e)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# checkpoint config-mismatch warnings (unchanged surface, still exercised)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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")
|
||||
ckpt_path.write_bytes(b"")
|
||||
_write_toml(
|
||||
tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n"
|
||||
)
|
||||
|
||||
gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
|
||||
|
||||
@@ -113,202 +659,9 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
|
||||
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")
|
||||
_write_toml(
|
||||
tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n"
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user