Scope wandb run config to only-active hyperparameters

Router-only knobs (lambda_balance/lambda_proc/lambda_entropy/
gumbel_tau_start/_end) and WGAN-only knobs (n_critic/gp_weight) were
being logged to wandb's top-level run config unconditionally, even for
runs where routing or WGAN mode is off, implying hyperparameters from
an inactive code path. Extract _wandb_run_config and only include each
group when its gate is actually true (router.enabled / mode=="wgan");
the full model_config (with its router sub-dict) is still always
logged in full, so no information is lost.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/train.py

@@ -126,6 +126,73 @@ def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -
     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,
@@ -412,29 +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,
-                "gumbel_tau_start": gumbel_tau_start,
-                "gumbel_tau_end": gumbel_tau_end,
-                "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)

tests/test_train.py

@@ -1,6 +1,6 @@
 """Tests for giant/train.py helpers."""
 
-from giant.train import _gumbel_tau
+from giant.train import _gumbel_tau, _wandb_run_config
 
 
 def test_gumbel_tau_at_step_zero_is_start():
@@ -24,3 +24,74 @@ def test_gumbel_tau_handles_zero_total_steps():
     # (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