From 8065df896e860248c49d55d61a4721975e6ed4c0 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 30 Jul 2026 16:46:19 +0200 Subject: [PATCH] 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 --- giant/train.py | 113 +++++++++++++++++++++++++++++++++++--------- tests/test_train.py | 73 +++++++++++++++++++++++++++- 2 files changed, 162 insertions(+), 24 deletions(-) diff --git a/giant/train.py b/giant/train.py index bcc8a38..0977a86 100644 --- a/giant/train.py +++ b/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) diff --git a/tests/test_train.py b/tests/test_train.py index 0b0a25d..627fe38 100644 --- a/tests/test_train.py +++ b/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