diff --git a/configs/router_energy_n10_gumbel_centers_physical.toml b/configs/router_energy_n10_gumbel_centers_physical.toml new file mode 100644 index 0000000..386e01e --- /dev/null +++ b/configs/router_energy_n10_gumbel_centers_physical.toml @@ -0,0 +1,22 @@ +[train] +mode = "flow" +epochs = 30 +lr = 3e-4 +warmup_epochs = 3 +val_fraction = 0.1 +num_workers = 4 + +[model] +conditioning = "physical" +dropout = 0.0 + +[model.router] +enabled = true +type = "energy" +n_experts = 10 +expert_hidden_dim = 128 +expert_n_blocks = 4 +temperature = 0.05 +lambda_balance = 0.035 +learn_centers = true +gumbel = true diff --git a/configs/router_energy_n10_gumbel_centers_temp_physical.toml b/configs/router_energy_n10_gumbel_centers_temp_physical.toml new file mode 100644 index 0000000..933aa4b --- /dev/null +++ b/configs/router_energy_n10_gumbel_centers_temp_physical.toml @@ -0,0 +1,23 @@ +[train] +mode = "flow" +epochs = 30 +lr = 3e-4 +warmup_epochs = 3 +val_fraction = 0.1 +num_workers = 4 + +[model] +conditioning = "physical" +dropout = 0.0 + +[model.router] +enabled = true +type = "energy" +n_experts = 10 +expert_hidden_dim = 128 +expert_n_blocks = 4 +temperature = 0.05 +lambda_balance = 0.035 +learn_centers = true +learn_temperature = true +gumbel = true diff --git a/configs/router_energy_n10_gumbel_physical.toml b/configs/router_energy_n10_gumbel_physical.toml new file mode 100644 index 0000000..51c6a39 --- /dev/null +++ b/configs/router_energy_n10_gumbel_physical.toml @@ -0,0 +1,22 @@ +[train] +mode = "flow" +epochs = 30 +lr = 3e-4 +warmup_epochs = 3 +val_fraction = 0.1 +num_workers = 4 + +[model] +conditioning = "physical" +dropout = 0.0 + +[model.router] +enabled = true +type = "energy" +n_experts = 10 +expert_hidden_dim = 128 +expert_n_blocks = 4 +temperature = 0.05 +lambda_balance = 0.035 +learn_centers = false +gumbel = true diff --git a/giant/config.py b/giant/config.py index d21be99..d021693 100644 --- a/giant/config.py +++ b/giant/config.py @@ -99,6 +99,19 @@ DEFAULT_CONFIG: dict = { # failure mode. Off by default; bounding above is the primary # defense. See giant.model.network.Router.entropy_loss. "lambda_entropy": 0.0, + # Opt-in straight-through Gumbel-softmax train-time combine weights + # (see giant.model.network.Router.combine_weights): the training + # forward pass samples a hard one-hot combination — matching + # eval-time top-1 dispatch exactly — while the backward pass still + # flows a smooth gradient to every expert. Targets the train/eval + # mismatch identified as a likely contributor to experts + # overlapping instead of partitioning (see CLAUDE.md roadmap). + # gumbel_tau_start/_end are annealed linearly over training + # (giant.train._gumbel_tau); off by default, no effect unless + # gumbel = true. + "gumbel": False, + "gumbel_tau_start": 1.0, + "gumbel_tau_end": 0.1, "emb_dim": 8, # process/pdg-router kwarg: own pdg(/mat) embedding width "hidden_dim": 64, # process-router kwarg: its classifier's hidden width "lambda_proc": 0.0, # process-router kwarg: supervised process-CE weight @@ -326,6 +339,29 @@ def _router_candidate(train, model): return f"r-{router['type']}{router['n_experts']}" +def _router_flag_candidate(field, token_map): + """Candidate factory for a boolean `model.router` sub-field. + + Gated on `router.enabled` like `_router_candidate` (a disabled router's + sub-fields are meaningless), then omitted unless `field` differs from + its DEFAULT_CONFIG value — same "only show non-default" rule as every + other candidate. `token_map` need only cover the non-default value(s), + since the default value always yields None. + """ + + def _candidate(train, model): + router = model["router"] + default_router = DEFAULT_CONFIG["model"]["router"] + if router["enabled"] == default_router["enabled"]: + return None + value = router[field] + if value == default_router[field]: + return None + return token_map[value] + + return _candidate + + def _conditioning_candidate(train, model): if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]: return None @@ -347,6 +383,10 @@ def _default_field_candidate(section_key, field, prefix): _OUT_DIR_NAME_CANDIDATES = [ ("mode", _mode_candidate), ("router", _router_candidate), + ("gumbel", _router_flag_candidate("gumbel", {True: "gum"})), + ("learn_centers", _router_flag_candidate("learn_centers", {False: "nolc"})), + ("learn_width", _router_flag_candidate("learn_width", {True: "lw"})), + ("learn_temperature", _router_flag_candidate("learn_temperature", {True: "lt"})), ("conditioning", _conditioning_candidate), ("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")), ("n_blocks", _default_field_candidate("model", "n_blocks", "b")), diff --git a/giant/model/network.py b/giant/model/network.py index f517d02..d04e4c4 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -537,11 +537,49 @@ class Router(nn.Module): def __init__(self, n_experts: int) -> None: super().__init__() self.n_experts = n_experts + # Opt-in straight-through Gumbel-softmax combine weights (see + # combine_weights below) — off by default, set from model.router.gumbel + # by _build_router_from_cfg. gumbel_tau is annealed per training step + # by giant.train (model.router.gumbel_tau_start/_end); neither is an + # nn.Parameter/buffer since neither is learned or needs checkpointing — + # the tau schedule is deterministic in global_step, so it recomputes + # correctly on resume. + self.gumbel = False + self.gumbel_tau = 1.0 def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: """(B, n_experts) soft weights, rows summing to 1.""" raise NotImplementedError + def combine_weights( + self, cond_cont: torch.Tensor, cond_cat: torch.Tensor + ) -> torch.Tensor: + """(B, n_experts) train-time expert-combination weights. + + Default (`gumbel=False`): identical to `gate()` — the original dense + soft-mixture combination. Opt-in straight-through Gumbel-softmax + (`gumbel=True`, train mode only): samples a Gumbel-perturbed + categorical draw from the same distribution `gate()` defines + (`log(gate())` is a valid unnormalized-logit input to + `F.gumbel_softmax` since softmax is shift-invariant, so no subclass + needs to expose separate pre-softmax logits), then hardens it to a + one-hot vector on the forward pass while keeping the soft sample's + gradient on the backward pass. This makes the training-time forward + combination match eval-time top-1 dispatch exactly (one expert's + output, unweighted) instead of the smooth blend `gate()` gives — + intended to close the train/eval mismatch identified as a likely + cause of experts overlapping instead of partitioning (see the + router_gating write-up referenced in CLAUDE.md's roadmap). + `gate()` itself is untouched and still backs `balance_loss`/ + `entropy_loss`/`gate_stats`, so those diagnostics keep reading the + smooth distribution rather than a noisy sample. + """ + probs = self.gate(cond_cont, cond_cat) + if not (self.gumbel and self.training): + return probs + log_probs = torch.log(probs.clamp_min(1e-8)) + return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1) + def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: """(B,) hard expert index, used for eval-time grouped dispatch.""" return self.gate(cond_cont, cond_cat).argmax(dim=-1) @@ -970,13 +1008,17 @@ def _route_forward( ) -> torch.Tensor: """Shared dispatch for both Routed* trunks. - Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully - differentiable, N-expert compute. Eval mode: grouped top-1 dispatch — - each row runs exactly one (small) expert, which is the actual source - of the per-call speedup this architecture is for. + Train mode: full mixture `sum_i weight_i * expert_i(x)` — always + N-expert dense compute, fully differentiable. `weight` is + `router.combine_weights(...)`: the plain soft `gate()` by default, or (see + `Router.combine_weights`) a straight-through Gumbel-softmax one-hot sample + when `router.gumbel` is enabled — either way, no change to the compute + cost of this branch. Eval mode: grouped top-1 dispatch — each row runs + exactly one (small) expert, which is the actual source of the per-call + speedup this architecture is for. """ if training: - weights = router.gate(cond_cont, cond_cat) # (B, n_experts) + weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts) out = torch.zeros_like(x) for i, expert in enumerate(experts): out = out + weights[:, i : i + 1] * expert(x, cond) @@ -1191,10 +1233,19 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair. + + `gumbel` is set as a post-construction attribute here rather than a + per-subclass constructor kwarg, same reasoning as `lambda_balance`/ + `lambda_proc`/`lambda_entropy` living in `router_cfg` without being a + `Router` subclass constructor param: it's a training-time toggle shared by + every router type, not a per-type hyperparameter (`build_router`'s + kwarg-filtering would otherwise just silently drop it). """ shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab) if router_cfg["type"] == "composed": - return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab) + router = build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab) + router.gumbel = bool(router_cfg.get("gumbel", False)) + return router router_kwargs = { k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts") } @@ -1204,7 +1255,9 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> # vocab, same as the trunk's ConditionEncoder. router_kwargs.setdefault("pdg_vocab", pdg_vocab) router_kwargs.setdefault("mat_vocab", mat_vocab) - return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs) + router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs) + router.gumbel = bool(router_cfg.get("gumbel", False)) + return router def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: diff --git a/giant/pipeline.py b/giant/pipeline.py index 53bdc61..ed12e55 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -417,6 +417,8 @@ def run_train_job( lambda_balance=router_cfg.get("lambda_balance", 0.0), lambda_proc=router_cfg.get("lambda_proc", 0.0), lambda_entropy=router_cfg.get("lambda_entropy", 0.0), + gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0), + gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1), normalizer_dict={ "cond": cond_norm.to_dict(), "target": tgt_norm.to_dict(), diff --git a/giant/train.py b/giant/train.py index b0d8654..0977a86 100644 --- a/giant/train.py +++ b/giant/train.py @@ -114,6 +114,85 @@ def _update_ema( ema_p.mul_(decay).add_(p, alpha=1 - decay) +def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -> float: + """Linear anneal of the straight-through Gumbel-softmax temperature. + + Deterministic in `step`/`total_steps` alone (no extra state), so it + recomputes correctly on `--resume` from a checkpoint's saved `global_step` + without needing to persist anything new (see + giant.model.network.Router.combine_weights). + """ + progress = min(step / max(total_steps, 1), 1.0) + return tau_start + (tau_end - tau_start) * progress + + +def _wandb_run_config( + *, + mode: str, + epochs: int, + lr: float, + warmup_epochs: int, + weight_decay: float, + ema_decay: float, + lambda_nsec: float, + lambda_s2: float, + lambda_balance: float, + lambda_proc: float, + lambda_entropy: float, + gumbel_tau_start: float, + gumbel_tau_end: float, + n_critic: int, + gp_weight: float, + model_config: dict | None, + stage1_params: int, + sec_decoder_params: int, + critic_params: int, + sec_critic_params: int, + total_params: int, +) -> dict: + """Build the dict logged as a wandb run's `config`. + + Router-only knobs (`lambda_balance`/`lambda_proc`/`lambda_entropy`/ + `gumbel_tau_start`/`gumbel_tau_end`) and WGAN-only knobs (`n_critic`/ + `gp_weight`) are omitted unless actually active, so a run's wandb config + doesn't imply hyperparameters from an inactive code path (a disabled + router's fine-tuning knobs, or GAN critic settings for a flow/DDPM run). + The full `model_config` (including its `router` sub-dict, whatever the + router type/state) is always included, so no information is lost — this + only trims the flattened top-level convenience duplicates. + """ + router_enabled = bool((model_config or {}).get("router", {}).get("enabled", False)) + cfg = { + "mode": mode, + "epochs": epochs, + "lr": lr, + "warmup_epochs": warmup_epochs, + "weight_decay": weight_decay, + "ema_decay": ema_decay, + "lambda_nsec": lambda_nsec, + "lambda_s2": lambda_s2, + "model": model_config or {}, + "stage1_params": stage1_params, + "sec_decoder_params": sec_decoder_params, + "critic_params": critic_params, + "sec_critic_params": sec_critic_params, + "total_params": total_params, + } + if router_enabled: + cfg.update( + { + "lambda_balance": lambda_balance, + "lambda_proc": lambda_proc, + "lambda_entropy": lambda_entropy, + "gumbel_tau_start": gumbel_tau_start, + "gumbel_tau_end": gumbel_tau_end, + } + ) + if mode == "wgan": + cfg.update({"n_critic": n_critic, "gp_weight": gp_weight}) + return cfg + + def _compute_losses( stage1_model: torch.nn.Module, sec_decoder: torch.nn.Module, @@ -347,6 +426,8 @@ def train( lambda_balance: float = 0.0, lambda_proc: float = 0.0, lambda_entropy: float = 0.0, + gumbel_tau_start: float = 1.0, + gumbel_tau_end: float = 0.1, normalizer_dict: dict | None = None, pdg_map: dict | None = None, mat_map: dict | None = None, @@ -398,27 +479,29 @@ def train( name=wandb_run_name or out_dir.name, id=out_dir.name, resume="allow", - config={ - "mode": mode, - "epochs": epochs, - "lr": lr, - "warmup_epochs": warmup_epochs, - "weight_decay": weight_decay, - "ema_decay": ema_decay, - "lambda_nsec": lambda_nsec, - "lambda_s2": lambda_s2, - "lambda_balance": lambda_balance, - "lambda_proc": lambda_proc, - "lambda_entropy": lambda_entropy, - "n_critic": n_critic, - "gp_weight": gp_weight, - "model": model_config or {}, - "stage1_params": stage1_params, - "sec_decoder_params": sec_decoder_params, - "critic_params": critic_params, - "sec_critic_params": sec_critic_params, - "total_params": total_params, - }, + config=_wandb_run_config( + mode=mode, + epochs=epochs, + lr=lr, + warmup_epochs=warmup_epochs, + weight_decay=weight_decay, + ema_decay=ema_decay, + lambda_nsec=lambda_nsec, + lambda_s2=lambda_s2, + lambda_balance=lambda_balance, + lambda_proc=lambda_proc, + lambda_entropy=lambda_entropy, + gumbel_tau_start=gumbel_tau_start, + gumbel_tau_end=gumbel_tau_end, + n_critic=n_critic, + gp_weight=gp_weight, + model_config=model_config, + stage1_params=stage1_params, + sec_decoder_params=sec_decoder_params, + critic_params=critic_params, + sec_critic_params=sec_critic_params, + total_params=total_params, + ), ) stage1_model = stage1_model.to(device) @@ -596,6 +679,13 @@ def train( dynamic_ncols=True, ) for batch in bar: + if has_router: + gumbel_tau = _gumbel_tau( + global_step, total_steps, gumbel_tau_start, gumbel_tau_end + ) + stage1_model.router.gumbel_tau = gumbel_tau + sec_decoder.router.gumbel_tau = gumbel_tau + if mode == "wgan": assert ( critic is not None @@ -737,6 +827,7 @@ def train( ) log_payload["batch/router_s1_entropy"] = s1_entropy.item() log_payload["batch/router_s2_entropy"] = s2_entropy.item() + log_payload["batch/gumbel_tau"] = gumbel_tau wandb_run.log(log_payload, step=global_step) if shutdown.requested: diff --git a/tests/test_config.py b/tests/test_config.py index 76f1da0..30f6e41 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -137,6 +137,16 @@ def test_resolve_expert_dims_missing_keys_also_inherit(): 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) @@ -197,6 +207,64 @@ def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefau 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" diff --git a/tests/test_router.py b/tests/test_router.py index 05105dc..34c669e 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -286,6 +286,110 @@ def test_router_entropy_loss_is_nonnegative_bounded_scalar(): assert 0.0 <= loss.item() <= 1.0 +# ── Router.combine_weights (straight-through Gumbel-softmax) ─────────────── + + +def test_combine_weights_defaults_to_gate(): + """gumbel=False (the default) must be a pure pass-through to gate().""" + router = EnergyRouter(n_experts=4) + cond_cont, cond_cat = _cond(16) + torch.testing.assert_close( + router.combine_weights(cond_cont, cond_cat), + router.gate(cond_cont, cond_cat), + ) + + +def test_combine_weights_gumbel_train_mode_is_hard_one_hot(): + router = EnergyRouter(n_experts=4) + router.gumbel = True + router.gumbel_tau = 0.5 + router.train() + cond_cont, cond_cat = _cond(16) + weights = router.combine_weights(cond_cont, cond_cat) + assert weights.shape == (16, 4) + torch.testing.assert_close(weights.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0) + assert torch.all((weights.max(dim=-1).values - 1.0).abs() < 1e-5) + + +def test_combine_weights_gumbel_eval_mode_falls_back_to_gate(): + """No Gumbel noise at eval — combine_weights must match gate() exactly, + same as the gumbel=False path, once the router is in eval mode.""" + router = EnergyRouter(n_experts=4) + router.gumbel = True + router.eval() + cond_cont, cond_cat = _cond(16) + torch.testing.assert_close( + router.combine_weights(cond_cont, cond_cat), + router.gate(cond_cont, cond_cat), + ) + + +def test_combine_weights_gumbel_straight_through_gradient_reaches_centers(): + router = EnergyRouter(n_experts=4, learn_centers=True) + router.gumbel = True + router.gumbel_tau = 0.5 + router.train() + cond_cont, cond_cat = _cond(16) + weights = router.combine_weights(cond_cont, cond_cat) + weights.sum().backward() + assert router.centers.grad is not None + assert torch.any(router.centers.grad != 0.0) + + +def test_build_router_from_cfg_sets_gumbel_from_config(): + from giant.model.network import _build_router_from_cfg + + router = _build_router_from_cfg( + {"enabled": True, "type": "energy", "n_experts": 4, "gumbel": True}, + pdg_vocab=3, + mat_vocab=2, + ) + assert router.gumbel is True + + router_off = _build_router_from_cfg( + {"enabled": True, "type": "energy", "n_experts": 4}, + pdg_vocab=3, + mat_vocab=2, + ) + assert router_off.gumbel is False + + +def test_build_router_from_cfg_sets_gumbel_for_composed_router(): + from giant.model.network import _build_router_from_cfg + + router = _build_router_from_cfg( + { + "enabled": True, + "type": "composed", + "gumbel": True, + "axis0_type": "energy", + "axis0_n_experts": 4, + "axis1_type": "pdg", + "axis1_n_experts": 3, + }, + pdg_vocab=5, + mat_vocab=2, + ) + assert isinstance(router, ComposedRouter) + assert router.gumbel is True + + +def test_routed_denoising_mlp_forward_runs_with_gumbel_enabled(): + """End-to-end forward through _route_forward's train branch with + straight-through Gumbel-softmax combine weights enabled.""" + B = 8 + model = _routed_stage1(n_experts=3) + model.router.gumbel = True + model.router.gumbel_tau = 0.5 + model.train() + x_t = torch.randn(B, X_DIM) + t = torch.rand(B) + cond_cont, cond_cat = _cond(B) + out = model(x_t, t, cond_cont, cond_cat) + assert out.shape == (B, X_DIM) + assert torch.isfinite(out).all() + + # ── PdgRouter ──────────────────────────────────────────────────────────────── diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 0000000..627fe38 --- /dev/null +++ b/tests/test_train.py @@ -0,0 +1,97 @@ +"""Tests for giant/train.py helpers.""" + +from giant.train import _gumbel_tau, _wandb_run_config + + +def test_gumbel_tau_at_step_zero_is_start(): + assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0 + + +def test_gumbel_tau_at_total_steps_is_end(): + assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9 + + +def test_gumbel_tau_interpolates_linearly_midway(): + assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9 + + +def test_gumbel_tau_clamps_beyond_total_steps(): + assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1) + + +def test_gumbel_tau_handles_zero_total_steps(): + # total_steps=0 is guarded to 1 internally: step=0 gives zero progress + # (still tau_start), any step>=1 immediately clamps to full progress. + assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0 + assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9 + + +def _base_wandb_kwargs(**overrides): + kwargs = dict( + mode="flow", + epochs=30, + lr=3e-4, + warmup_epochs=3, + weight_decay=0.01, + ema_decay=0.9999, + lambda_nsec=0.1, + lambda_s2=1.0, + lambda_balance=0.035, + lambda_proc=0.0, + lambda_entropy=0.0, + gumbel_tau_start=1.0, + gumbel_tau_end=0.1, + n_critic=5, + gp_weight=10.0, + model_config={"router": {"enabled": False}}, + stage1_params=100, + sec_decoder_params=50, + critic_params=0, + sec_critic_params=0, + total_params=150, + ) + kwargs.update(overrides) + return kwargs + + +def test_wandb_run_config_omits_router_knobs_when_router_disabled(): + cfg = _wandb_run_config(**_base_wandb_kwargs()) + for key in ( + "lambda_balance", + "lambda_proc", + "lambda_entropy", + "gumbel_tau_start", + "gumbel_tau_end", + ): + assert key not in cfg + # still present, nested, regardless of router state + assert cfg["model"] == {"router": {"enabled": False}} + + +def test_wandb_run_config_includes_router_knobs_when_router_enabled(): + cfg = _wandb_run_config( + **_base_wandb_kwargs(model_config={"router": {"enabled": True}}) + ) + assert cfg["lambda_balance"] == 0.035 + assert cfg["lambda_proc"] == 0.0 + assert cfg["lambda_entropy"] == 0.0 + assert cfg["gumbel_tau_start"] == 1.0 + assert cfg["gumbel_tau_end"] == 0.1 + + +def test_wandb_run_config_omits_wgan_knobs_when_mode_is_not_wgan(): + cfg = _wandb_run_config(**_base_wandb_kwargs(mode="flow")) + assert "n_critic" not in cfg + assert "gp_weight" not in cfg + + +def test_wandb_run_config_includes_wgan_knobs_when_mode_is_wgan(): + cfg = _wandb_run_config(**_base_wandb_kwargs(mode="wgan")) + assert cfg["n_critic"] == 5 + assert cfg["gp_weight"] == 10.0 + + +def test_wandb_run_config_handles_missing_model_config(): + cfg = _wandb_run_config(**_base_wandb_kwargs(model_config=None)) + assert cfg["model"] == {} + assert "lambda_balance" not in cfg